rocm: disable gradient_accumulation_fusion on gfx950 in test_sglang_config#1156
rocm: disable gradient_accumulation_fusion on gfx950 in test_sglang_config#1156sreerohi wants to merge 2 commits into
Conversation
hipBLASLt on gfx950 (MI350/MI355) has no algorithm for the triple combination of bf16 output with fp32 accumulate + HIPBLASLT_EPILOGUE_BGRADB epilogue + accumulate=True. This fires during TE's LayerNormLinear backward when gradient_accumulation_fusion=True and the layer has bias. Root cause: the bridge provider defaults gradient_accumulation_fusion to True (via can_enable_gradient_accumulation_fusion) and ignores --no-gradient-accumulation-fusion from the CLI. The fix propagates the flag from Megatron args to the bridge provider.
sreerohi
requested review from
Zhichenzzz,
fzyzcjy,
maocheng23,
yueming-yuan and
yushengsu-thu
as code owners
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request enables the gradient_accumulation_fusion flag to be controlled via Megatron CLI arguments to address compatibility issues on ROCm/gfx950. It also updates the E2E configuration tests to conditionally disable specific checkers and fusions when running in a ROCm environment. Review feedback highlights a potential AttributeError when checking the PyTorch version on non-ROCm systems and suggests using a consistent null-check pattern for configuration arguments to avoid overwriting default values.
|
|
||
| import miles.utils.external_utils.command_utils as U | ||
|
|
||
| IS_ROCM = torch.version.hip is not None |
There was a problem hiding this comment.
Accessing torch.version.hip directly will raise an AttributeError on non-ROCm (e.g., CUDA) builds of PyTorch, as the hip attribute is only defined in ROCm-enabled builds. This will cause the test to crash during collection or execution on CUDA environments. Use getattr to safely check for its existence.
| IS_ROCM = torch.version.hip is not None | |
| IS_ROCM = getattr(torch.version, "hip", None) is not None |
| provider.gradient_accumulation_fusion = getattr( | ||
| args, "gradient_accumulation_fusion", provider.gradient_accumulation_fusion | ||
| ) |
There was a problem hiding this comment.
To maintain consistency with the existing code in this file (e.g., lines 105-112) and to avoid overwriting the provider's default value if the argument is None, consider using the if ... is not None pattern. This ensures that the provider's default configuration is preserved unless explicitly overridden by a non-null CLI argument.
if getattr(args, "gradient_accumulation_fusion", None) is not None:
provider.gradient_accumulation_fusion = args.gradient_accumulation_fusionReferences
- Maintain consistency with existing patterns in the file for handling CLI arguments and defaults. (link)
1 participant
Depends on #1153. Related to #1105 (Miles CI gap between ROCm & CUDA).
hipBLASLt on gfx950 has no algorithm for TE's bias-fused wgrad GEMM (bf16→fp32 + BGRADB + accumulate). This conditionally disables gradient_accumulation_fusion and relaxes CI numerical checkers on ROCm. CUDA is unaffected.