Describe the bug
When training the qwen3-vl-8b model, the generic_gemm operator of the reference throws an addition violation error when performing the out = out + bias operation. The incorrect shape is "A.shape=torch.Size([960, 1152]), B.shape=torch.Size([960, 1728]), transA=False, transB=True, bias.shape=torch.Size([1728])". The reason may be that in line 873 of transformer_engine/pytorch/module/layernorm_linear.py, the bias is passed incorrectly. The code on line 873 is ""bias": (bias if (grad_bias is None and not ctx.fp8) else None),".
Steps/Code to reproduce bug
The steps I followed based on the README file at https://github.com/flagos-ai/FlagScale/blob/main/examples/qwen3_vl/README.md were almost incorrect. I nearly forgot to install Megatron-LM-FL and instead set the environment variable for the Megatron-LM-FL path, skipping the checkpoint conversion. I also used the demo dataset, and pretrained_checkpoint wasn't set. However, I don't think these near misses were the main factors contributing to the error.
Expected behavior
When the model training reaches the shape mentioned earlier, the input bias should be either none or the correct 1152. This should allow the training to succeed.
Environment overview (please complete the following information)
- Environment location: docker
- Method of Transformer Engine install: The te-fl directory was placed under site-packages, but pip install was not performed.
- If method of install is [Docker], provide
docker pull & docker run commands used：Docker uses a simplified environment, which I downloaded myself, including Flagscale, Megatron-lm-fl, and Te-fl.
Environment details
If NVIDIA docker image is used you don't need to specify these.
Otherwise, please provide:
- OS version:Ubuntu 20.04.6 LTS
- PyTorch version:2.5.1+cu118
- Python version:Python 3.10.16
- Transformer Engine version:1.13.0+4725d23d
- CUDA version:cuda_11.7
- CUDNN version:90100
Device details
Additional context
Add any other context about the problem here.
Describe the bug
When training the qwen3-vl-8b model, the generic_gemm operator of the reference throws an addition violation error when performing the out = out + bias operation. The incorrect shape is "A.shape=torch.Size([960, 1152]), B.shape=torch.Size([960, 1728]), transA=False, transB=True, bias.shape=torch.Size([1728])". The reason may be that in line 873 of transformer_engine/pytorch/module/layernorm_linear.py, the bias is passed incorrectly. The code on line 873 is ""bias": (bias if (grad_bias is None and not ctx.fp8) else None),".
Steps/Code to reproduce bug
The steps I followed based on the README file at https://github.com/flagos-ai/FlagScale/blob/main/examples/qwen3_vl/README.md were almost incorrect. I nearly forgot to install Megatron-LM-FL and instead set the environment variable for the Megatron-LM-FL path, skipping the checkpoint conversion. I also used the demo dataset, and
pretrained_checkpointwasn't set. However, I don't think these near misses were the main factors contributing to the error.Expected behavior
When the model training reaches the shape mentioned earlier, the input bias should be either none or the correct 1152. This should allow the training to succeed.
Environment overview (please complete the following information)
docker pull&docker runcommands used：Docker uses a simplified environment, which I downloaded myself, including Flagscale, Megatron-lm-fl, and Te-fl.Environment details
If NVIDIA docker image is used you don't need to specify these.
Otherwise, please provide:
Device details
Additional context
Add any other context about the problem here.