Major cleanup and refactor

Get rid of legacy names like packed offloading
Move the main code body of paged stash to transformer/moe/

Author: nanz-nv

megatron/core/full_cuda_graph.py

@@ -7,8 +7,8 @@
 import torch
 
 from megatron.core.tensor_parallel.random import get_all_rng_states
-from megatron.core.pipeline_parallel.moe_packed_offload import (
-    packed_moe_expert_offloading_reset,
+from megatron.core.transformer.moe.paged_stash import (
+    paged_stash_reset,
 )
 
 logger = logging.getLogger(__name__)
@@ -101,11 +101,11 @@ class FullCudaGraphWrapper:
     cuda_graph = {'training': None, 'validation': None}
     result = {'training': None, 'validation': None}
 
-    def __init__(self, forward_backward_func, cuda_graph_warmup_steps=1, packed_moe_expert_offloading=False):
+    def __init__(self, forward_backward_func, cuda_graph_warmup_steps=1, moe_paged_stash=False):
         self.forward_backward_func = forward_backward_func
         self.static_loader = StaticBufferLoader()
         self.cuda_graph_warmup_steps = cuda_graph_warmup_steps
-        self.packed_moe_expert_offloading = packed_moe_expert_offloading
+        self.moe_paged_stash = moe_paged_stash
 
     def data_read(self, data_iterator, model, training, num_microbatches):
         """Read all microbatch inputs from Dataloader and copy to static buffers."""
@@ -188,15 +188,15 @@ def __call__(self, *args, **kwargs):
         if FullCudaGraphWrapper.cuda_graph[training_str] is None:
             FullCudaGraphWrapper.result[training_str] = self.forward_backward_func(*args, **kwargs)
         else:
-            packed_moe_expert_offloading_reset(enabled=self.packed_moe_expert_offloading and training)
+            paged_stash_reset(enabled=self.moe_paged_stash and training)
             FullCudaGraphWrapper.cuda_graph[training_str].replay()
         self.speculative_cuda_graph_check(model)
         self.next_iter(training_str)
         return FullCudaGraphWrapper.result[training_str]
 
     def speculative_cuda_graph_check(self, model):
         ''' check speculative execution modules '''
-        if self.packed_moe_expert_offloading:
+        if self.moe_paged_stash:
             # Check if there is any overflow in the receiving buffer
             over_budget = torch.zeros(1, dtype=torch.bool, device='cuda')
             for model_chunk in model:

megatron/core/models/common/model_chunk_schedule_plan.py

@@ -11,8 +11,8 @@
 from megatron.core.pipeline_parallel.fine_grained_activation_offload import (
     fine_grained_offloading_set_last_layer,
 )
-from megatron.core.pipeline_parallel.moe_packed_offload import (
-    packed_moe_expert_offloading_set_last_layer,
+from megatron.core.transformer.moe.paged_stash import (
+    paged_stash_set_last_layer,
 )
 from megatron.core.pipeline_parallel.utils import (
     AbstractSchedulePlan,
@@ -460,8 +460,8 @@ def run(
             torch.cuda.nvtx.range_push(f"layer_{i}f-layer_{b_num_layers - 1 - i}b")
             if f_layer.layer.config.fine_grained_activation_offloading:
                 fine_grained_offloading_set_last_layer(i == f_num_layers - 1)
-            if f_layer.layer.config.packed_moe_expert_offloading:
-                packed_moe_expert_offloading_set_last_layer(i == f_num_layers - 1)
+            if f_layer.layer.config.moe_paged_stash:
+                paged_stash_set_last_layer(i == f_num_layers - 1)
             f_input, b_grad = TransformerLayerSchedulePlan.run(
                 f_layer,
                 b_layer,
@@ -486,8 +486,8 @@ def run(
             torch.cuda.nvtx.range_push(f"layer_{i}f")
             if f_layer.layer.config.fine_grained_activation_offloading:
                 fine_grained_offloading_set_last_layer(i == f_num_layers - 1)
-            if f_layer.layer.config.packed_moe_expert_offloading:
-                packed_moe_expert_offloading_set_last_layer(i == f_num_layers - 1)
+            if f_layer.layer.config.moe_paged_stash:
+                paged_stash_set_last_layer(i == f_num_layers - 1)
             f_input, _ = TransformerLayerSchedulePlan.run(f_layer, None, f_input=f_input)
             torch.cuda.nvtx.range_pop()
 

megatron/core/models/gpt/gpt_model.py

@@ -21,8 +21,8 @@
 from megatron.core.pipeline_parallel.fine_grained_activation_offload import (
     fine_grained_offloading_init_chunk_handler,
 )
-from megatron.core.pipeline_parallel.moe_packed_offload import (
-    packed_moe_expert_offloading_init_chunk_handler,
+from megatron.core.transformer.moe.paged_stash import (
+    paged_stash_init_chunk_handler,
 )
 from megatron.core.process_groups_config import ProcessGroupCollection
 from megatron.core.quantization.utils import get_quant_config_or_none
@@ -435,9 +435,9 @@ def preprocess_for_fine_grained_offloading(self):
                     param.offloading_activation = False
             self.disable_param_offloading = False
 
-    def preprocess_for_packed_moe_expert_offloading(self):
-        """Preprocess for packed moe expert offloading."""
-        return packed_moe_expert_offloading_init_chunk_handler(
+    def preprocess_for_paged_stash(self):
+        """Preprocess for paged stash."""
+        return paged_stash_init_chunk_handler(
             vp_size=self.config.virtual_pipeline_model_parallel_size,
             vp_stage=self.vp_stage,
         )
@@ -470,8 +470,8 @@ def forward(
         if self.config.fine_grained_activation_offloading:
             self.preprocess_for_fine_grained_offloading()
 
-        if self.config.packed_moe_expert_offloading:
-            self.preprocess_for_packed_moe_expert_offloading()
+        if self.config.moe_paged_stash:
+            self.preprocess_for_paged_stash()
 
         inference_context = deprecate_inference_params(inference_context, inference_params)
 
@@ -770,8 +770,8 @@ def build_schedule_plan(
 
         if self.config.fine_grained_activation_offloading:
             self.preprocess_for_fine_grained_offloading()
-        if self.config.packed_moe_expert_offloading:
-            self.preprocess_for_packed_moe_expert_offloading()
+        if self.config.moe_paged_stash:
+            self.preprocess_for_paged_stash()
 
         from ..common.model_chunk_schedule_plan import TransformerModelChunkSchedulePlan