From f56c0b780eacfbcb09b882eb7c3bd2cf33735a26 Mon Sep 17 00:00:00 2001 From: hongbinl Date: Tue, 12 May 2026 02:18:24 -0700 Subject: [PATCH] [Dev] fix: restore PR #3219 fine-grained offload semantics after PR #4291 sync MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The main→dev sync in #4291 partially reverted #3219: 1. bulk_offload_group() regained pre-#3219 semantics (`group_to_offload = self._groups_to_offload[-1]` + `.pop()`), silently overriding the parameter and making bulk_offload()'s find_group_with_name() result effectively dead. Restore #3219's intent: honor the passed group, and let the caller remove it by identity via .remove(group_to_offload). 2. multi_latent_attention.py core_attn block was reverted to the old `with off_interface(...)` + `off_interface.group_commit(...)` API. Refactor to #3219's manager pattern: `core_attn_manager = off_interface(...)` once, then `with core_attn_manager as query:` and an unconditional `core_attn_manager.group_offload(...)` (no-op when offload flag is False). Matches the sibling qkv_linear / attn_proj blocks. 3. Drop the FineGrainedActivationOffloadingInterface.group_commit static-method shim added during the sync — it existed only to keep MLA's reverted call site working and has no remaining callers after (2). Signed-off-by: Hongbin Liu --- .../fine_grained_activation_offload.py | 10 +--------- .../core/transformer/multi_latent_attention.py | 14 +++++++------- 2 files changed, 8 insertions(+), 16 deletions(-) diff --git a/megatron/core/pipeline_parallel/fine_grained_activation_offload.py b/megatron/core/pipeline_parallel/fine_grained_activation_offload.py index 1e3601aef26..1901a3c07da 100644 --- a/megatron/core/pipeline_parallel/fine_grained_activation_offload.py +++ b/megatron/core/pipeline_parallel/fine_grained_activation_offload.py @@ -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): @@ -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 @@ -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() @@ -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.""" diff --git a/megatron/core/transformer/multi_latent_attention.py b/megatron/core/transformer/multi_latent_attention.py index c8f00084c5d..6954f39a7fe 100644 --- a/megatron/core/transformer/multi_latent_attention.py +++ b/megatron/core/transformer/multi_latent_attention.py @@ -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 @@ -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, @@ -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():