Skip to content

[ROCm] Fix biased wgrad with fp32 gradient accumulation #634

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
XinyuJiangCMU wants to merge 5 commits into
base: dev
Choose a base branch
from
+16 −12
Open

[ROCm] Fix biased wgrad with fp32 gradient accumulation #634

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
38d375d
[ROCm] Skip fused bias-grad in wgrad GEMM on fp32-accumulate path
XinyuJiangCMU Jun 11, 2026
c9904e3
Hoist ROCm wgrad bias-grad unfuse to general_gemm
XinyuJiangCMU Jun 15, 2026
1715351
Trim verbose comments in wgrad bias-grad split
XinyuJiangCMU Jun 17, 2026
627e0ef
[ROCm] Trigger wgrad bias-grad split on fp32 output, not accumulate
XinyuJiangCMU Jun 18, 2026
abe3697
[ROCm] Trigger wgrad bias-grad split on fp32 output, not accumulate
XinyuJiangCMU Jun 18, 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
11 changes: 0 additions & 11 deletions tests/pytorch/test_numerics.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1410,10 +1410,6 @@ def test_linear_accuracy_delay_wgrad_compute(dtype, bs, model, bias, fuse_wgrad_

config = model_configs[model]

if IS_HIP_EXTENSION:
if dtype not in (torch.float32,) and fuse_wgrad_accumulation and bias:
pytest.skip(f"ROCm does not support fused wgrad accumulation for {dtype}.")

te_linear_ref = Linear(
config.hidden_size,
4 * config.hidden_size,
Expand Down Expand Up @@ -1708,9 +1704,6 @@ def test_layernorm_linear_accuracy(
def test_layernorm_linear_accuracy_delay_wgrad_compute(
dtype, bs, model, normalization, zero_centered_gamma, bias, fuse_wgrad_accumulation
):
if IS_HIP_EXTENSION:
if dtype not in (torch.float32,) and fuse_wgrad_accumulation and bias:
pytest.skip(f"ROCm does not support fused wgrad accumulation for {dtype}.")
if NVTE_TEST_NVINSPECT_ENABLED:
pytest.skip("Delayed wgrad compute is not supported in debug mode.")

Expand Down Expand Up @@ -1932,10 +1925,6 @@ def test_layernorm_mlp_accuracy_delay_wgrad_compute(

config = model_configs[model]

if IS_HIP_EXTENSION:
if dtype not in (torch.float32,) and fuse_wgrad_accumulation and bias:
pytest.skip(f"ROCm does not support fused wgrad accumulation for {dtype}.")

ln_mlp = LayerNormMLP(
hidden_size=config.hidden_size,
ffn_hidden_size=4 * config.hidden_size,
Expand Down
17 changes: 16 additions & 1 deletion transformer_engine/pytorch/cpp_extensions/gemm.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -431,6 +431,18 @@ def general_gemm(
)
out_dtype = torch.bfloat16 if use_bf16_tn_output_workaround else out_dtype

# hipBLASLt has no fused bias-grad (BGRADB) algorithm for an fp32-output wgrad GEMM, so skip the fused dbias and reduce grad_output below.
rocm_split_dbias = (
IS_HIP_EXTENSION
and grad
and bias is not None
and not gelu
and out is not None
and out.dtype == torch.float32
and quantization_params is None
)
gemm_bias = None if rocm_split_dbias else bias

args = (
A,
transa, # transa
Expand All @@ -439,7 +451,7 @@ def general_gemm(
out,
quantization_params,
TE_DType[out_dtype] if out_dtype is not None else None,
bias,
gemm_bias,
bias_dtype,
gelu,
gelu_in,
Expand All @@ -460,6 +472,9 @@ def general_gemm(

out, bias_grad, gelu_input, extra_output = tex.generic_gemm(*args, **kwargs)

if rocm_split_dbias:
bias_grad = B.reshape(-1, B.shape[-1]).sum(dim=0, dtype=torch.float32).to(bias.dtype)

if IS_HIP_EXTENSION and use_bf16_tn_output_workaround:
out = cast_if_needed(out, torch.float32)

Expand Down
Loading