[Megatron-LM] feat: add grouped gemm fp4 support, skip cache trans weight and santize code by RuibinCheung · Pull Request #891 · AMD-AGI/Primus

RuibinCheung · 2026-07-20T10:09:54Z

Description

This PR adds FP4 (MXFP4) grouped GEMM support to PrimusTurboGroupedLinear.

It also fixes memory footprint by skipping the parameter transpose cache under current scaling.

Finally, it removes the experimental OPT-1 fused grouped-wgrad path together with its debug scaffolding to clean up the code.

Fixes # (issue)

Type of change

Documentation change (change only to the documentation, either a fix or a new content)
Bug fix (non-breaking change which fixes an issue)
New feature (non-breaking change which adds functionality)
Breaking change (fix or feature that would cause existing functionality to not work as expected)
Infra/Build change
Code refactoring

Changes

Please list the changes introduced in this PR:

Add FP4 grouped GEMM support to PrimusTurboGroupedLinear.
When Turbo FP4 is enabled, build MXFP4 (float4_e2m1fn_x2) quantized weight buffers and dispatch to grouped_gemm_fp4, replacing the previous assert False, "FP4 is not supported in PrimusTurboGroupedLinear".
Skip caching the transposed weight when current scaling is active.
disable_parameter_transpose_cache is now OR'ed with quant_config.current_scaling() across PrimusTurboLinear, PrimusTurboRowParallelLinear, PrimusTurboColumnParallelLinear, PrimusTurboLayerNormColumnParallelLinear, and PrimusTurboGroupedLinear, so a stale transposed weight is not reused under current scaling.
Remove the experimental OPT-1 fused grouped-wgrad path and its scaffolding.
Drop the _MainGradShim helper, the fused_grouped_wgrad / _expert_main_grad_view integration, the PRIMUS_TURBO_FUSE_GROUPED_WGRAD / PRIMUS_TURBO_FUSE_WGRAD_DEBUG env-var gating and debug prints, and the associated is_gfx1250 Triton in-place add workaround in _bridge_weight_grad.
Simplify _bridge_weight_grad to always accumulate grad_quantized_weight into weight.main_grad and flag grad_added_to_main_grad.
Remove the now-unused PRIMUS_TURBO_FUSE_GROUPED_WGRAD and PRIMUS_TURBO_FUSE_WGRAD_DEBUG env vars from examples/deepseek-v4/run_deepseek_v4_pro_muon_1gpu.sh, and drop the unused contextlib / os imports.

Checklist:

The functionality is complete
I have commented my code, particularly in hard-to-understand areas
I have made corresponding changes to the documentation
My changes generate no new warnings
I have added tests that prove my fix is effective or that my feature works
New and existing unit tests pass locally with my changes

Copilot

Pull request overview

This PR simplifies the Megatron Primus Turbo wrapper by removing the experimental OPT-1 “fused grouped wgrad” single-GPU path and its associated debug/env-var scaffolding, returning PrimusTurboGroupedLinear to the standard grouped_gemm_fp8 backward flow.

Changes:

Removed the OPT-1 fused grouped-wgrad branch (and _MainGradShim) from PrimusTurboGroupedLinear, always using the plain grouped_gemm_fp8 call.
Simplified _bridge_weight_grad to always accumulate grad_quantized_weight into weight.main_grad, dropping the old “fused path returns None” branch and the gfx1250 Triton add workaround.
Stopped forwarding PRIMUS_TURBO_FUSE_GROUPED_WGRAD / PRIMUS_TURBO_FUSE_WGRAD_DEBUG into the DeepSeek V4 1GPU example run script.

Reviewed changes

Copilot reviewed 2 out of 2 changed files in this pull request and generated 1 comment.

File	Description
primus/backends/megatron/core/extensions/primus_turbo.py	Removes the OPT-1 fused wgrad/debug path and simplifies FP8 weight-grad bridging back to the default grouped GEMM behavior.
examples/deepseek-v4/run_deepseek_v4_pro_muon_1gpu.sh	Drops forwarding of the removed OPT-1 env vars to the container runtime.

+            weight.main_grad.add_(grad_quantized_weight)
+            weight.grad_added_to_main_grad = True


…ze code

Copilot

Pull request overview

Copilot reviewed 2 out of 2 changed files in this pull request and generated 1 comment.

+        elif PrimusTurboLowPrecisionGlobalStateManager.is_turbo_fp4_enabled():
+            quant_config = PrimusTurboLowPrecisionGlobalStateManager.get_turbo_quant_config()
+            assert quant_config.mxfp4_scaling(), "Turbo FP4 is enabled but quant config is not mxfp4."



cursor

Stale comment

安全审查结论：未发现此 PR 引入或暴露的中危及以上高置信漏洞。新增 FP4 grouped GEMM 路径仅接收框架内部张量和受限量化配置；同时删除了环境变量驱动的实验性梯度路径，未扩大攻击面。因此本轮不提交行内安全问题。

_{Sent by Cursor Automation: Find vulnerabilities}

cursor

Stale comment

安全审查结论：未发现此 PR 引入或暴露的中危及以上高置信漏洞。新增 grouped FP4 路径仅消费框架内部生成的量化配置、权重和 MoE token 计数，并沿用现有 grouped GEMM 信任边界；配置限定为 MXFP4 且进入上下文时会检查硬件支持。其余改动删除了环境变量驱动的实验性梯度分支，未扩大攻击面。因此本轮不提交行内安全问题。

_{Sent by Cursor Automation: Find vulnerabilities}

Copilot

Pull request overview

Copilot reviewed 2 out of 2 changed files in this pull request and generated no new comments.

Comments suppressed due to low confidence (1)

primus/backends/megatron/core/extensions/primus_turbo.py:190

_WeightGradBridge.backward unconditionally calls weight.main_grad.add_(grad_quantized_weight). If autograd supplies None for grad_quantized_weight (e.g., if an upstream op returns no wgrad), this will raise a TypeError with a hard-to-debug message. Add an explicit assertion (or guard) so failures are clearer and intentional.

            weight.main_grad.add_(grad_quantized_weight)
            weight.grad_added_to_main_grad = True

cursor

Stale comment

安全审查结论：未发现此 PR 当前版本引入或暴露的中危及以上高置信漏洞。新增 MXFP4 grouped GEMM 路径仅接收框架内部张量、由路由结果生成的 token 计数以及受 MXFP4 配方和硬件能力检查约束的结构化量化配置；未发现攻击者可控输入通向注入、越权、秘密泄漏、SSRF、路径遍历或不安全反序列化 sink。其余改动删除了环境变量驱动的实验性梯度分支并禁用 current-scaling 权重转置缓存，未扩大安全边界。因此本轮不提交行内安全问题。

_{Sent by Cursor Automation: Find vulnerabilities}

 try:
-    import triton
-    import triton.language as tl
+    pass


Copilot

Pull request overview

Copilot reviewed 6 out of 6 changed files in this pull request and generated 1 comment.

Comments suppressed due to low confidence (1)

primus/backends/megatron/core/extensions/primus_turbo.py:175

_WeightGradBridge.backward assumes grad_quantized_weight is always a Tensor. Previously this code handled grad_quantized_weight is None (e.g., when a fused/accumulate-in-main-grad path returns no explicit weight grad). Without a guard, this will raise when grad_quantized_weight is None. Also, if Triton is unavailable, the gfx1250 branch should fall back to add_() even if _is_gfx1250() returns true.

            if _is_gfx1250():
                inplace_add_triton_(weight.main_grad, grad_quantized_weight)
            else:
                weight.main_grad.add_(grad_quantized_weight)
            weight.grad_added_to_main_grad = True

 try:
-    import triton
-    import triton.language as tl
+    pass

    _HAVE_TRITON = True
 except (ImportError, ModuleNotFoundError):
    _HAVE_TRITON = False


cursor

Stale comment

安全审查结论：未发现此 PR 当前版本引入或暴露的中危及以上高置信漏洞，因此未提交新的行内 finding。已重点复核最新提交新增的 gfx1250 Triton 梯度累加内核、单次 microbatch 的 FP8/FP4 未缓存量化路径及 grouped MXFP4 GEMM：内核访问由实际目标张量 numel 和掩码限制，输入均来自同一训练进程内的模型参数、autograd 梯度或由 MoE 路由结果生成且受 token 数量约束的计数；量化配置由受限 recipe 构造并经过格式及硬件能力检查。未发现攻击者可控输入通向注入、越权、秘密泄漏、SSRF、路径遍历、不安全反序列化或可利用越界读写 sink。此前无本自动化留下的未解决 finding 线程；旧评估已清理。

_{Sent by Cursor Automation: Find vulnerabilities}

Copilot

Pull request overview

Copilot reviewed 6 out of 6 changed files in this pull request and generated 1 comment.

Comments suppressed due to low confidence (2)

primus/backends/megatron/core/extensions/primus_turbo.py:94

The Triton availability guard is broken: the try block contains only pass, so _HAVE_TRITON is always set to True, and the subsequent unconditional import of inplace_add_triton_ will still raise ImportError when Triton isn’t installed. This both misreports capability and removes the previous graceful fallback behavior.

try:
    pass

    _HAVE_TRITON = True
except (ImportError, ModuleNotFoundError):

primus/backends/megatron/core/extensions/primus_turbo.py:175

_bridge_weight_grad assumes inplace_add_triton_ is always importable. If Triton (or the helper module) isn’t available, this will crash on gfx1250 even though there is a safe Torch fallback.

            if _is_gfx1250():
                inplace_add_triton_(weight.main_grad, grad_quantized_weight)
            else:
                weight.main_grad.add_(grad_quantized_weight)
            weight.grad_added_to_main_grad = True

+    if not torch.cuda.is_available():
+        return False
+    try:
+        return "gfx1250" in torch.cuda.get_device_properties(0).gcnArchName
+    except Exception:
+        return False


cursor

Stale comment

安全审查结论：未发现此 PR 当前版本引入或暴露的中危及以上高置信漏洞，因此未提交新的行内 finding。已复核此前评估及最新同步；当前合并提交相对前一 PR 头未改变代码。新增 Triton inplace_add 的访问范围由目标张量 numel 与逐元素掩码限制，调用输入来自同一训练进程内的模型参数和 autograd 梯度；新增 grouped MXFP4 路径接收框架生成的 MoE token 计数及受 MXFP4 recipe/硬件能力检查约束的量化配置。未发现攻击者可控输入通向注入、越权、秘密泄漏、SSRF、路径遍历、不安全反序列化或可利用越界读写 sink。此前本自动化没有未解决的安全 finding 线程；旧评估已清理。

_{Sent by Cursor Automation: Find vulnerabilities}

Copilot

Pull request overview

Copilot reviewed 6 out of 6 changed files in this pull request and generated 1 comment.

Comments suppressed due to low confidence (2)

primus/backends/megatron/core/extensions/primus_turbo.py:94

The Triton availability guard is currently a no-op (try: pass) so _HAVE_TRITON is always set to True, and inplace_add_triton_ is imported unconditionally. This makes the module fail to import in environments where Triton isn’t installed, and the _HAVE_TRITON flag is misleading.

Consider restoring a real optional import that both (1) detects Triton correctly and (2) avoids importing the Triton helper when Triton is unavailable.

try:
    pass

    _HAVE_TRITON = True
except (ImportError, ModuleNotFoundError):

primus/backends/megatron/core/extensions/primus_turbo.py:175

inplace_add_triton_ can be unavailable if Triton isn’t installed (or if the import is intentionally made optional). The current code calls it unconditionally on gfx1250, which would raise at runtime in that case. Guard the call and fall back to Tensor.add_ when the Triton helper isn’t available.

            if _is_gfx1250():
                inplace_add_triton_(weight.main_grad, grad_quantized_weight)
            else:
                weight.main_grad.add_(grad_quantized_weight)
            weight.grad_added_to_main_grad = True

+    if not dst.is_cuda or not dst.is_contiguous() or dst.numel() != src.numel():
+        return dst.add_(src)


cursor

Stale comment

安全审查结论：未发现此 PR 当前版本引入或暴露的中危及以上高置信漏洞，因此不提交新的行内 finding。

已复核历史审查线程：此前本自动化没有未解决的安全 finding。当前同步仅将 main 合入分支，PR 修改的 6 个文件与上一轮已审查头提交逐字节一致。新增 Triton 梯度累加内核的访问上界由目标张量 numel 和逐元素掩码约束，调用参数来自同一训练进程内的模型参数与 autograd 梯度；grouped MXFP4 路径仅消费框架生成的张量、MoE token 计数及通过 MXFP4 类型和硬件能力检查的量化配置。未发现攻击者可控输入可到达注入、越权、秘密泄漏、SSRF、路径遍历、不安全反序列化或可利用的越界内存访问 sink。旧评估已清理。

_{Sent by Cursor Automation: Find vulnerabilities}

Copilot

Pull request overview

Copilot reviewed 6 out of 6 changed files in this pull request and generated no new comments.

Comments suppressed due to low confidence (3)

primus/backends/megatron/core/extensions/primus_turbo.py:94

The try/except block here is effectively a no-op (pass) and will always set _HAVE_TRITON = True, even when Triton is not installed. It also makes the subsequent Triton import unconditional, so the module can fail to import in environments without Triton. Replace this with a real optional import that defines inplace_add_triton_ when available and sets it to None otherwise (and drop _HAVE_TRITON since it’s unused).

try:
    pass

    _HAVE_TRITON = True
except (ImportError, ModuleNotFoundError):

primus/backends/megatron/core/extensions/primus_turbo.py:116

_is_gfx1250() queries device properties for device index 0, which can be the wrong GPU in multi-GPU jobs (each rank typically sets a non-zero current device). Use torch.cuda.current_device() (and getattr for gcnArchName) to avoid mis-detecting the architecture.

    if not torch.cuda.is_available():
        return False
    try:
        return "gfx1250" in torch.cuda.get_device_properties(0).gcnArchName
    except Exception:

primus/backends/megatron/core/extensions/primus_turbo.py:175

grad_quantized_weight can legally be None in PyTorch autograd (e.g., if the corresponding input doesn’t require grad or an upstream op returns None). The current code unconditionally does weight.main_grad.add_(grad_quantized_weight), which will throw a TypeError if grad_quantized_weight is None. Guard against None and also only call inplace_add_triton_ when it’s available (it should be None when Triton isn’t installed).

            if _is_gfx1250():
                inplace_add_triton_(weight.main_grad, grad_quantized_weight)
            else:
                weight.main_grad.add_(grad_quantized_weight)
            weight.grad_added_to_main_grad = True

Copilot

Pull request overview

Copilot reviewed 6 out of 6 changed files in this pull request and generated no new comments.

Comments suppressed due to low confidence (2)

primus/backends/megatron/core/extensions/primus_turbo.py:99

The Triton availability probe is currently a no-op (try: pass) and sets _HAVE_TRITON = True unconditionally. _HAVE_TRITON is also unused, while inplace_add_triton_ is imported unconditionally and will raise at import time if Triton isn't installed. This makes the module behavior inconsistent and can break environments that previously relied on the Torch fallback.

try:
    pass

    _HAVE_TRITON = True
except (ImportError, ModuleNotFoundError):
    _HAVE_TRITON = False

from primus.backends.megatron.core.extensions._triton.inplace_add import (
    inplace_add_triton_,
)

primus/backends/megatron/core/extensions/primus_turbo.py:116

_is_gfx1250() queries torch.cuda.get_device_properties(0), which can be the wrong device on multi-GPU nodes (or when the current device is not 0). This can incorrectly enable/disable the gfx1250 workaround on a given rank.

    if not torch.cuda.is_available():
        return False
    try:
        return "gfx1250" in torch.cuda.get_device_properties(0).gcnArchName
    except Exception:

cursor

安全审查结论：未发现此 PR 当前版本引入或暴露的中危及以上高置信漏洞，因此不提交新的行内 finding。

已复核历史审查线程；此前本自动化没有未解决的安全 finding。新增 Triton 梯度累加内核的访问范围受目标张量 numel 和逐元素掩码约束，输入来自同一训练进程内的模型参数与 autograd 梯度；新增 grouped MXFP4 路径仅消费路由生成且受 token 数约束的 m_splits、框架内部张量，以及通过 MXFP4 类型和硬件能力检查的结构化量化配置。其余改动删除实验性环境变量路径或仅调整缓存策略/文档，未扩大安全边界。旧评估已清理。

_{Sent by Cursor Automation: Find vulnerabilities}

Copilot AI review requested due to automatic review settings July 20, 2026 10:09

RuibinCheung requested review from Xiaoming-AMD, limou102 and wenxie-amd as code owners July 20, 2026 10:09

Copilot started reviewing on behalf of RuibinCheung July 20, 2026 10:10 View session

Copilot AI reviewed Jul 20, 2026

View reviewed changes

Comment thread primus/backends/megatron/core/extensions/primus_turbo.py Outdated

Comment on lines +189 to +190

weight.main_grad.add_(grad_quantized_weight)

weight.grad_added_to_main_grad = True

feat: add grouped gemm fp4 support, skip cache trans weight and santi…

27f8c9c

…ze code

Copilot AI review requested due to automatic review settings July 22, 2026 07:20

RuibinCheung force-pushed the fix/zhangrb/clean_code branch from 59a49d7 to 27f8c9c Compare July 22, 2026 07:20

RuibinCheung changed the title ~~[Megatron-LM] fix: santize primus turbo wrapper code~~ [Megatron-LM] feat: add grouped gemm fp4 support, skip cache trans weight and santize code Jul 22, 2026

Copilot started reviewing on behalf of RuibinCheung July 22, 2026 07:21 View session

Copilot AI reviewed Jul 22, 2026

View reviewed changes

cursor Bot reviewed Jul 22, 2026

View reviewed changes

Merge branch 'main' into fix/zhangrb/clean_code

062cf0c

cursor Bot reviewed Jul 22, 2026

View reviewed changes

Merge branch 'main' into fix/zhangrb/clean_code

7219b55

Copilot AI review requested due to automatic review settings July 23, 2026 01:54

Copilot started reviewing on behalf of RuibinCheung July 23, 2026 01:55 View session

Copilot AI reviewed Jul 23, 2026

View reviewed changes

cursor Bot reviewed Jul 23, 2026

View reviewed changes

skip quantized weight when micro batch is 1 and clean some code

88a7019

Copilot AI review requested due to automatic review settings July 23, 2026 04:15

Copilot started reviewing on behalf of RuibinCheung July 23, 2026 04:17 View session

github-code-quality Bot found potential problems Jul 23, 2026

View reviewed changes

Comment thread primus/backends/megatron/core/extensions/primus_turbo.py

try:

import triton

import triton.language as tl

pass

Copilot AI reviewed Jul 23, 2026

View reviewed changes

Comment thread primus/backends/megatron/core/extensions/primus_turbo.py

Comment on lines 90 to 95

try:

import triton

import triton.language as tl

pass

_HAVE_TRITON = True

except (ImportError, ModuleNotFoundError):

_HAVE_TRITON = False

cursor Bot reviewed Jul 23, 2026

View reviewed changes

Merge branch 'main' into fix/zhangrb/clean_code

49abb96

Copilot AI review requested due to automatic review settings July 23, 2026 08:53

Copilot started reviewing on behalf of wenxie-amd July 23, 2026 08:54 View session

Copilot AI reviewed Jul 23, 2026

View reviewed changes

Comment thread primus/backends/megatron/core/extensions/primus_turbo.py

Comment on lines +112 to +117

if not torch.cuda.is_available():

return False

try:

return "gfx1250" in torch.cuda.get_device_properties(0).gcnArchName

except Exception:

return False

Copilot AI review requested due to automatic review settings July 23, 2026 08:59

Copilot started reviewing on behalf of RuibinCheung July 23, 2026 09:00 View session

cursor Bot reviewed Jul 23, 2026

View reviewed changes

Copilot AI reviewed Jul 23, 2026

View reviewed changes

Comment thread primus/backends/megatron/core/extensions/_triton/inplace_add.py

Comment on lines +48 to +49

if not dst.is_cuda or not dst.is_contiguous() or dst.numel() != src.numel():

return dst.add_(src)

Merge branch 'main' into fix/zhangrb/clean_code

03027c9

Copilot AI review requested due to automatic review settings July 23, 2026 09:06

Copilot started reviewing on behalf of RuibinCheung July 23, 2026 09:06 View session

cursor Bot reviewed Jul 23, 2026

View reviewed changes

Copilot AI reviewed Jul 23, 2026

View reviewed changes

Merge branch 'main' into fix/zhangrb/clean_code

bca3142

Copilot AI review requested due to automatic review settings July 23, 2026 11:05

Copilot started reviewing on behalf of RuibinCheung July 23, 2026 11:06 View session

Copilot AI reviewed Jul 23, 2026

View reviewed changes

cursor Bot reviewed Jul 23, 2026

View reviewed changes

		weight.main_grad.add_(grad_quantized_weight)
		weight.grad_added_to_main_grad = True

		if not dst.is_cuda or not dst.is_contiguous() or dst.numel() != src.numel():
		return dst.add_(src)

Uh oh!

Conversation

RuibinCheung commented Jul 20, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Type of change

Changes

Checklist:

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Uh oh!

cursor Bot left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

cursor Bot left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Uh oh!

cursor Bot left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Uh oh!

cursor Bot left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Uh oh!

cursor Bot left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Uh oh!

cursor Bot left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Uh oh!

cursor Bot left a comment

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

RuibinCheung commented Jul 20, 2026 •

edited

Loading

cursor Bot left a comment •

edited

Loading

cursor Bot left a comment •

edited

Loading

cursor Bot left a comment •

edited

Loading

cursor Bot left a comment •

edited

Loading

cursor Bot left a comment •

edited

Loading

cursor Bot left a comment •

edited

Loading