Skip to content

[Dev] fix: restore PR #3219 fine-grained offload semantics after PR #… #4757

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

Merged
lhb8125 merged 2 commits into from
May 13, 2026
Merged

[Dev] fix: restore PR #3219 fine-grained offload semantics after PR #… #4757

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
f56c0b7
[Dev] fix: restore PR #3219 fine-grained offload semantics after PR #…
lhb8125 May 12, 2026
c5fdbde
Merge branch 'dev' into hongbinl/fix_3219_bulk_offload_group
lhb8125 May 12, 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
10 changes: 1 addition & 9 deletions megatron/core/pipeline_parallel/fine_grained_activation_offload.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -957,7 +957,6 @@ def tensor_need_offloading_checker(self, tensor):
def bulk_offload_group(self, group_to_offload):
"""offload a group of tensors recorded in tensor_push()."""
debug_rank("------bulk_offload_group")
group_to_offload = self._groups_to_offload[-1]
nvtx_msg = "activation offloading " + group_to_offload._name
nvtx_range_push(nvtx_msg)
with torch.cuda.stream(self.d2h_stream):
Expand All @@ -971,7 +970,6 @@ def bulk_offload_group(self, group_to_offload):
tensor_on_device.record_stream(self.d2h_stream)
group_to_offload.push_tensor(tensor_tag, state)
group_to_offload.record_offload_event(self.d2h_stream)
self._groups_to_offload.pop()
nvtx_range_pop(nvtx_msg)
# Under full-iteration CG capture, the main stream may not wait on d2h
# events; optional max-inflight enqueues each group's offload event and
Expand Down Expand Up @@ -1055,6 +1053,7 @@ def bulk_offload(self, name, forced_released_tensors):
), f"Group {name} not found in {self._groups_to_offload}"
self._groups_to_reload.append(group_to_offload)
self.bulk_offload_group(group_to_offload)
self._groups_to_offload.remove(group_to_offload)
# Manually release tensors not auto-freed by torch GC
if len(forced_released_tensors) > 0:
cur_stream = torch.cuda.current_stream()
Expand Down Expand Up @@ -1370,13 +1369,6 @@ def group_offload(self, tensor, forced_released_tensors=None, delay_offload=Fals
)
return tensor

@staticmethod
def group_commit(tensor, name, forced_released_tensors=None, delay_offload=False):
"""Static variant of group_offload used by main's multi_latent_attention."""
return fine_grained_offloading_group_commit(
tensor, name, forced_released_tensors, delay_offload
)

@staticmethod
def mark_not_offload(tensor: torch.Tensor):
"""Mark the tensor as not offloadable."""
Expand Down
14 changes: 7 additions & 7 deletions megatron/core/transformer/multi_latent_attention.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -368,6 +368,9 @@ def forward(
# ==================================
# Need corresponding TE change
needs_output_trim = False
core_attn_manager = off_interface(
self.offload_core_attention and self.training, query, "core_attn"
)
if self.checkpoint_core_attention and self.training:
core_attn_out = self._checkpointed_attention_forward(
query, key, value, attention_mask, packed_seq_params=packed_seq_params
Expand All @@ -380,9 +383,7 @@ def forward(
# query representation.
extra_kwargs["x"] = hidden_states
extra_kwargs["qr"] = q_compressed
with off_interface(
self.offload_core_attention and self.training, query, "core_attn"
) as query:
with core_attn_manager as query:
core_attn_out = self._run_core_attention(
query,
key,
Expand Down Expand Up @@ -416,10 +417,9 @@ def forward(
if not inference_context.is_decode_only():
core_attn_out = rearrange(core_attn_out, 's b h d -> s b (h d)')
needs_output_trim = need_v_pad
if self.offload_core_attention and self.training:
core_attn_out = off_interface.group_commit(
core_attn_out, name="core_attn", forced_released_tensors=[query, key, value]
)
core_attn_out = core_attn_manager.group_offload(
core_attn_out, forced_released_tensors=[query, key, value]
)

# We are doing absorption with cache mla latents and decode mode.
if self.cache_mla_latents and inference_context.is_decode_only():
Expand Down
Loading