diff --git a/miles/backends/megatron_utils/model_provider.py b/miles/backends/megatron_utils/model_provider.py index 6aa694e876..5ab27401ee 100644 --- a/miles/backends/megatron_utils/model_provider.py +++ b/miles/backends/megatron_utils/model_provider.py @@ -95,6 +95,107 @@ def _apply_bridge_runtime_config(provider, args: argparse.Namespace) -> None: provider.dsa_attention_backend = getattr(args, "dsa_attention_backend", "megatron") +_ROCM_TE_BGRADB_WRAPPER_INSTALLED = False +def _install_rocm_te_bgradb_fallback() -> None: + """Wrap TE's `tex.generic_gemm` so that the bias-fused wgrad GEMM + (HIPBLASLT_EPILOGUE_BGRADB) transparently falls back to the unfused + path when hipBLASLt has no algorithm for the (M, N, K) being issued. + + On ROCm/gfx950 (MI350-355) hipBLASLt's algorithm catalog has holes in + the BGRADB epilogue for some shapes (e.g. Qwen2.5-0.5B's QKV wgrad: + BF16, M ~ max_tokens_per_gpu, N=1152, K=896), raising + `RuntimeError: Unable to find any suitable algorithms` from + `tex.generic_gemm`. We catch that specific error when the call is a + bias-fused wgrad (bias != None, grad == True), retry the GEMM with + bias=None (default epilogue, which has working algorithms for every + shape we've seen on gfx950), and compute the bias gradient ourselves + via `grad_output.sum(dim=0)`. Mathematically identical to BGRADB; just + not fused into the same kernel. + + This keeps the full TE module stack (LayerNormLinear, fused norm+linear + forward, etc.) and only diverges from TE's behaviour for the specific + catalog miss, so there are no model-parameter-name changes and no perf + cost when hipBLASLt has the algorithm. + + The wrapper is a no-op on non-ROCm platforms. + """ + global _ROCM_TE_BGRADB_WRAPPER_INSTALLED + if _ROCM_TE_BGRADB_WRAPPER_INSTALLED: + return + if torch.version.hip is None: + _ROCM_TE_BGRADB_WRAPPER_INSTALLED = True + return + + try: + import transformer_engine_torch as tex + except ImportError: + logger.warning( + "transformer_engine_torch not importable; skipping ROCm BGRADB " + "fallback install." + ) + _ROCM_TE_BGRADB_WRAPPER_INSTALLED = True + return + + if getattr(tex, "_miles_rocm_bgradb_wrapped", False): + _ROCM_TE_BGRADB_WRAPPER_INSTALLED = True + return + + original_generic_gemm = tex.generic_gemm + + # Index layout of positional args to tex.generic_gemm, mirrored from + # transformer_engine/pytorch/cpp_extensions/gemm.py::general_gemm: + # 0: A, 1: transa, 2: B, 3: transb, 4: out, 5: quantization_params, + # 6: out_dtype, 7: bias, 8: bias_dtype, 9: gelu, 10: gelu_in, + # 11: grad, 12: workspace, 13: workspace_size, 14: accumulate, + # 15: use_split_accumulator + _BIAS_IDX = 7 + _GRAD_IDX = 11 + _B_IDX = 2 + + def wrapped_generic_gemm(*args, **kwargs): + try: + return original_generic_gemm(*args, **kwargs) + except RuntimeError as e: + if "Unable to find any suitable algorithms" not in str(e): + raise + # Only fall back when this is the bias-fused wgrad case. + if len(args) <= _GRAD_IDX: + raise + bias = args[_BIAS_IDX] + grad = args[_GRAD_IDX] + if bias is None or not grad: + raise + logger.warning( + "hipBLASLt has no algorithm for the bias-fused wgrad GEMM " + "(HIPBLASLT_EPILOGUE_BGRADB) at this shape; retrying with " + "bias=None and computing bias_grad via sum(dim=0)." + ) + # Retry the GEMM without the bias-grad epilogue. + new_args = list(args) + new_args[_BIAS_IDX] = None + out, _, gelu_input, extra_output = original_generic_gemm(*new_args, **kwargs) + # TE issues this call with layout="NT" for wgrad: B is the + # grad_output, shape (M, N). The BGRADB epilogue would have + # computed bias_grad = B.sum over M -> (N,). Replicate that. + B = args[_B_IDX] + bias_dtype = bias.dtype if hasattr(bias, "dtype") else B.dtype + bias_grad = B.sum(dim=0).to(dtype=bias_dtype) + return out, bias_grad, gelu_input, extra_output + + tex.generic_gemm = wrapped_generic_gemm + tex._miles_rocm_bgradb_wrapped = True + _ROCM_TE_BGRADB_WRAPPER_INSTALLED = True + logger.info( + "Installed ROCm BGRADB fallback wrapper around " + "transformer_engine_torch.generic_gemm." + ) + + +# Install the ROCm BGRADB fallback as soon as this module is imported, before +# any model is built so that TE backwards passes go through the wrapper. +_install_rocm_te_bgradb_fallback() + + # Adapt from https://github.com/volcengine/verl/blob/c3b20575d2bc815fcccd84bddb4c0401fc4b632b/verl/models/llama/megatron/layers/parallel_linear.py#L82 class LinearForLastLayer(torch.nn.Linear): def __init__(