Skip to content

[fix] Respect per-op activation-offload markers in fused grouped MLP #3128

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

Closed
fanshiqing wants to merge 1 commit into from
Closed

[fix] Respect per-op activation-offload markers in fused grouped MLP #3128

Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
29 changes: 27 additions & 2 deletions transformer_engine/pytorch/ops/fused/grouped_mlp.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,12 @@

import transformer_engine_torch as tex
from ...constants import MXFP8_BLOCK_SCALING_SIZE, NVFP4_BLOCK_SCALING_SIZE
from ...cpu_offload import is_cpu_offload_enabled, mark_activation_offload, start_offload
from ...cpu_offload import (
is_cpu_offload_enabled,
mark_activation_offload,
mark_not_offload,
start_offload,
)
from ...cpp_extensions import general_gemm, general_grouped_gemm_for_grouped_tensor
from ...module.base import _2X_ACC_WGRAD
from ...quantization import Recipe
Expand Down Expand Up @@ -1450,9 +1455,29 @@ def fuser_forward(
grouped_fc_x.rowwise_data = None
grouped_fc_x.scale_inv = None

# Per-op fine-grained offload markers.
offload_fc1_x = bool(getattr(fc1_op, "fine_grained_activation_offloading", False))
offload_act = bool(getattr(activation_op, "fine_grained_activation_offloading", False))
fine_grained_offload = offload_fc1_x or offload_act
saved_activations = (
(grouped_fc1_x, offload_fc1_x),
(activation_in, offload_act),
(saved_grouped_fc2_x, offload_act),
)

# The hook-based offloader is opt-out, so explicitly keep the
# non-selected tensors resident (mark_not_offload sets _TE_do_not_offload).
if fine_grained_offload:
keep = [t for t, sel in saved_activations if t is not None and not sel]
if keep:
mark_not_offload(*keep)
Comment on lines +1470 to +1473

@greptile-apps greptile-apps Bot Jun 14, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

P2 mark_not_offload called when no offload context is active

mark_not_offload is called whenever fine_grained_offload=True, even when cpu_offloading=False (i.e., no offloading context has been entered). In that case the hook (push_tensor) is never installed, so _TE_do_not_offload is never checked. More importantly, for the QuantizedTensor types used here (grouped_fc1_x, activation_in, saved_grouped_fc2_x), mark_not_offload internally calls prepare_for_saving, which is a destructive-then-restore operation — unnecessary work when no offloader is active. Guard the block with if cpu_offloading and fine_grained_offload: to match the start_offload / mark_activation_offload scope below.


if cpu_offloading:
# TE-native path; with no markers, offload everything saved (legacy).
activation_tensors = [
t for t in (grouped_fc1_x, activation_in, saved_grouped_fc2_x) if t is not None
t
for t, sel in saved_activations
if t is not None and (sel or not fine_grained_offload)
]
start_offload(*activation_tensors)
mark_activation_offload(*activation_tensors)
Expand Down
Loading