Skip to content

rocm: disable gradient_accumulation_fusion on gfx950 in test_sglang_config_mixed_offload #1157

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 2 commits into
base: main
Choose a base branch
from
Open

rocm: disable gradient_accumulation_fusion on gfx950 in test_sglang_config_mixed_offload #1157

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
f86c2f2
rocm: fix hipBLASLt BGRADB wgrad failure on gfx950
sreerohi May 19, 2026
18f011a
rocm: gate BGRADB workaround flags on IS_ROCM in test_sglang_config_m…
sreerohi May 19, 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
9 changes: 9 additions & 0 deletions miles/backends/megatron_utils/model_provider.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -110,6 +110,15 @@ def wrapped_model_provider(
provider.moe_router_bias_update_rate = args.moe_router_bias_update_rate
if getattr(args, "moe_aux_loss_coeff", None) is not None:
provider.moe_aux_loss_coeff = args.moe_aux_loss_coeff
# 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 on lines +119 to +121
Copy link
Copy Markdown
Contributor

@gemini-code-assist gemini-code-assist Bot May 19, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

To maintain consistency with the surrounding code (e.g., lines 105-112) and to avoid potentially overwriting the provider's default with None if the attribute exists in args but is uninitialized, it is recommended to use an explicit is not None check before assignment.

        if getattr(args, "gradient_accumulation_fusion", None) is not None:
            provider.gradient_accumulation_fusion = args.gradient_accumulation_fusion

provider.finalize()

def wrapped_bridge_provider(
Expand Down
11 changes: 9 additions & 2 deletions tests/e2e/sglang_config/test_sglang_config_mixed_offload.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,10 +16,13 @@
import os
import tempfile

import torch
from tests.ci.ci_register import register_cuda_ci

import miles.utils.external_utils.command_utils as U

IS_ROCM = torch.version.hip is not None

register_cuda_ci(est_time=600, suite="stage-b-short-8-gpu", num_gpus=8)

TIGHT_DEVICE_MEMORY = U.get_bool_env_var("MILES_TEST_TIGHT_DEVICE_MEMORY", "1")
Expand Down Expand Up @@ -119,13 +122,17 @@ def execute():
f"--sglang-config {config_path} "
)

ci_args = "--ci-test "
ci_args = (
"--ci-test "
+ ("--ci-disable-kl-checker --ci-disable-logprobs-checker " if IS_ROCM else "")
)

misc_args = (
"--attention-dropout 0.0 "
"--hidden-dropout 0.0 "
"--accumulate-allreduce-grads-in-fp32 "
"--attention-softmax-in-fp32 "
+ ("--no-gradient-accumulation-fusion " if IS_ROCM else "")
+ "--attention-softmax-in-fp32 "
"--attention-backend flash "
"--actor-num-nodes 1 "
"--actor-num-gpus-per-node 8 "
Expand Down