Skip to content

rocm: fix hipBLASLt BGRADB wgrad failure on gfx950 #1618

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
sreerohi wants to merge 4 commits into
base: main
Choose a base branch
from
+15 −0
Open

rocm: fix hipBLASLt BGRADB wgrad failure on gfx950 #1618

Changes from 1 commit
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
5b808ae
rocm: fix hipBLASLt BGRADB wgrad failure on gfx950
sreerohi May 19, 2026
a0b9499
Update miles/backends/megatron_utils/model_provider.py
sreerohi Jul 13, 2026
742e4d6
Add shared IS_ROCM detection util for CI tests
sreerohi Jul 14, 2026
e74b943
Update model_provider.py
sreerohi Jul 14, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 9 additions & 1 deletion miles/backends/megatron_utils/model_provider.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -93,7 +93,15 @@ def _apply_bridge_runtime_config(provider, args: argparse.Namespace) -> None:

if hasattr(provider, "dsa_attention_backend"):
provider.dsa_attention_backend = getattr(args, "dsa_attention_backend", "megatron")

# The bridge provider defaults gradient_accumulation_fusion=True
# (via fusions.can_enable_gradient_accumulation_fusion). On ROCm/gfx950,
# this makes TE's LayerNormLinear backward issue a bias-fused wgrad GEMM
# with bf16 inputs → fp32 output + HIPBLASLT_EPILOGUE_BGRADB + accumulate,
# for which hipBLASLt has no algorithm. Honor the Megatron CLI flag so
# that --no-gradient-accumulation-fusion actually takes effect.
provider.gradient_accumulation_fusion = getattr(
args, "gradient_accumulation_fusion", provider.gradient_accumulation_fusion
)
Comment thread
sreerohi marked this conversation as resolved.
Outdated

# Adapt from https://github.com/volcengine/verl/blob/c3b20575d2bc815fcccd84bddb4c0401fc4b632b/verl/models/llama/megatron/layers/parallel_linear.py#L82
class LinearForLastLayer(torch.nn.Linear):
Expand Down