diff --git a/miles/backends/megatron_utils/model_provider.py b/miles/backends/megatron_utils/model_provider.py index 6aa694e876..ed88f53553 100644 --- a/miles/backends/megatron_utils/model_provider.py +++ b/miles/backends/megatron_utils/model_provider.py @@ -93,6 +93,16 @@ 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. + if hasattr(provider, "gradient_accumulation_fusion"): + provider.gradient_accumulation_fusion = getattr( + args, "gradient_accumulation_fusion", provider.gradient_accumulation_fusion + ) # Adapt from https://github.com/volcengine/verl/blob/c3b20575d2bc815fcccd84bddb4c0401fc4b632b/verl/models/llama/megatron/layers/parallel_linear.py#L82 diff --git a/tests/ci/rocm_utils.py b/tests/ci/rocm_utils.py new file mode 100644 index 0000000000..c2e6d2a86f --- /dev/null +++ b/tests/ci/rocm_utils.py @@ -0,0 +1,5 @@ +"""Shared ROCm detection for CI tests.""" + +import torch + +IS_ROCM = getattr(torch.version, "hip", None) is not None