verl-project
diff --git a/‎verl/utils/veomni_utils.py‎
Lines changed: 0 additions & 78 deletions b/‎verl/utils/veomni_utils.py‎
Lines changed: 0 additions & 78 deletions
diff --git a/‎verl/workers/engine/veomni/transformer_impl.py‎
Lines changed: 126 additions & 43 deletions b/‎verl/workers/engine/veomni/transformer_impl.py‎
Lines changed: 126 additions & 43 deletions
@@ -15,7 +15,7 @@
 
 import logging
 from dataclasses import dataclass, field
-from typing import Any, Callable, Sequence
+from typing import Any, Callable, Optional, Sequence
 
 import torch
 import torch.distributed as dist
@@ -33,20 +33,22 @@
 from verl.utils.checkpoint.fsdp_checkpoint_manager import FSDPCheckpointManager
 from verl.utils.device import get_device_id, get_device_name
 from verl.utils.fsdp_utils import fsdp_version
+from verl.utils.model import convert_weight_keys
 from verl.utils.profiler import log_gpu_memory_usage
-from verl.utils.veomni_utils import (
-    load_veomni_model_to_gpu,
-    load_veomni_optimizer,
-    offload_veomni_model_to_cpu,
-    offload_veomni_optimizer,
-)
 from verl.workers.config import HFModelConfig, VeOmniEngineConfig, VeOmniOptimizerConfig
 from verl.workers.sharding_manager.fsdp_ulysses import FSDPUlyssesShardingManager
 
 from ..base import BaseEngineCtx, EngineRegistry
 from ..fsdp.transformer_impl import FSDPEngine, FSDPEngineWithLMHead
 from ..utils import enable_full_determinism, postprocess_batch_func, prepare_micro_batches
-from .utils import VL_TYPE2INDEX
+from .utils import (
+    MOE_PARAM_HANDERS,
+    VL_TYPE2INDEX,
+    load_veomni_model_to_gpu,
+    load_veomni_optimizer,
+    offload_veomni_model_to_cpu,
+    offload_veomni_optimizer,
+)
 
 logger = logging.getLogger(__file__)
 
@@ -61,23 +63,19 @@ def __init__(
         **kwargs,
     ):
         """
-        Initialize the FSDPEngine.
+        Initialize the VeOmniEngine.
 
         Sets up distributed device meshes, LoRA, and offload policies based on config.
 
         Args:
-            config: Configuration object with FSDP and model settings.
+            config: Configuration object with VeOmni and model settings.
         """
 
-        # TODO: Preprocessing operations for the MOE model are appended here,
-        # instead of relying on Veomni's transformation scripts.
-
         self.model_config = model_config
         self.engine_config = engine_config
         self.optimizer_config = optimizer_config
         self.checkpoint_config = checkpoint_config
-
-        self.mode = None
+        assert self.engine_config.data_parallel_mode == "fsdp2", "VeOmniEngine only supports fsdp2."
 
         self.rank = dist.get_rank()
 
@@ -223,34 +221,6 @@ def _build_model_optimizer(self):
             self.engine_config.activation_gpu_limit,
         )
 
-    def to(self, device: str, model: bool = True, optimizer: bool = True, grad: bool = True):
-        """
-        Move model parameters, optimizer states, or both to the specified device.
-        Note that this function executes irrespective of offload config. It serves as manual control.
-
-        Args:
-            device: Target device identifier.
-            model: If True, move the model.
-            optimizer: If True, move the optimizer states.
-        """
-        super(FSDPEngine, self).to(device=device, model=model, optimizer=optimizer, grad=grad)
-
-        device_name = get_device_name()
-
-        assert device in (device_name, "cpu")
-        if device == device_name:
-            if model:
-                load_veomni_model_to_gpu(self.module)
-            if optimizer and self.optimizer is not None:
-                load_veomni_optimizer(self.optimizer, device)
-        elif device == "cpu":
-            if model:
-                offload_veomni_model_to_cpu(self.module)
-            if optimizer and self.optimizer is not None:
-                offload_veomni_optimizer(self.optimizer)
-        else:
-            raise ValueError(f"Invalid device type: {device}")
-
     def optimizer_step(self):
         """
         Perform an optimization step using the optimizer.
@@ -348,6 +318,117 @@ def eval_mode(self, **kwargs):
         """
         return EngineEvalModeCtx(self, **kwargs)
 
+    def to(self, device: str, model: bool = True, optimizer: bool = True, grad: bool = True):
+        """
+        Move model parameters, optimizer states, or both to the specified device.
+        Note that this function executes irrespective of offload config. It serves as manual control.
+
+        Args:
+            device: Target device identifier.
+            model: If True, move the model.
+            optimizer: If True, move the optimizer states.
+        """
+        super(FSDPEngine, self).to(device=device, model=model, optimizer=optimizer, grad=grad)
+
+        device_name = get_device_name()
+
+        assert device in (device_name, "cpu")
+        if device == device_name:
+            if model:
+                load_veomni_model_to_gpu(self.module)
+            if optimizer and self.optimizer is not None:
+                load_veomni_optimizer(self.optimizer, device)
+        elif device == "cpu":
+            if model:
+                offload_veomni_model_to_cpu(self.module)
+            if optimizer and self.optimizer is not None:
+                offload_veomni_optimizer(self.optimizer)
+        else:
+            raise ValueError(f"Invalid device type: {device}")
+
+    def save_checkpoint(
+        self,
+        local_path: str,
+        hdfs_path: Optional[str] = None,
+        global_step: int = 0,
+        max_ckpt_to_keep: Optional[int] = None,
+        **kwargs,
+    ) -> None:
+        """
+        Save VeOmni checkpoint, handling parameter offload as needed.
+        """
+        origin_module_device = next(self.module.parameters()).device.type
+        if self._is_offload_param or origin_module_device == "cpu":
+            load_veomni_model_to_gpu(self.module)
+
+        self.checkpoint_manager.save_checkpoint(
+            local_path=local_path, hdfs_path=hdfs_path, global_step=global_step, max_ckpt_to_keep=max_ckpt_to_keep
+        )
+
+        torch.distributed.barrier()
+        if self._is_offload_param:
+            offload_veomni_model_to_cpu(self.module)
+
+    def load_checkpoint(
+        self, local_path: str, hdfs_path: Optional[str] = None, del_local_after_load: int = True, **kwargs
+    ) -> None:
+        """
+        Load VeOmni checkpoint, restoring parameters and optimizer state.
+        """
+        if self._is_offload_param:
+            load_veomni_model_to_gpu(self.module)
+
+        self.checkpoint_manager.load_checkpoint(
+            local_path=local_path, hdfs_path=hdfs_path, del_local_after_load=del_local_after_load
+        )
+
+        torch.distributed.barrier()
+        if self._is_offload_param:
+            offload_veomni_model_to_cpu(self.module)
+
+        if self._is_offload_optimizer:
+            offload_veomni_optimizer(self.optimizer)
+
+    def get_per_tensor_param(self, **kwargs):
+        load_veomni_model_to_gpu(self.module)
+
+        params = self.module.state_dict()
+        params = convert_weight_keys(params, getattr(self.module, "_fsdp_wrapped_module", self.module))
+
+        if self._is_offload_param:
+            offload_veomni_model_to_cpu(self.module)
+
+        device = get_device_id()
+        ps = parallel_state.get_parallel_state()
+        model_type = getattr(self.module.config, "model_type", "default")
+        process_func = MOE_PARAM_HANDERS.get(model_type, lambda n, t: iter([(n, t)]))
+
+        def param_generator():
+            for name, param in params.items():
+                unsharded_tensor = param.full_tensor() if isinstance(param, DTensor) else param
+
+                is_expert_layer = "mlp.experts." in name
+                is_proj = any(p in name for p in ["down_proj", "gate_proj", "up_proj", "gate_up_proj"])
+
+                if is_expert_layer and is_proj and ps.ep_enabled:
+                    output_shape = list(unsharded_tensor.shape)
+                    output_shape[0] *= ps.ep_size
+                    stacked_tensor = torch.empty(output_shape, dtype=unsharded_tensor.dtype, device=device)
+
+                    # all gather expert tensors [32, H, I] -> [128, H, I]
+                    torch.distributed.all_gather_into_tensor(stacked_tensor, unsharded_tensor, group=ps.ep_group)
+                    yield from process_func(name, stacked_tensor)
+
+                    del stacked_tensor
+                else:
+                    if is_expert_layer:
+                        yield from process_func(name, unsharded_tensor)
+                    else:
+                        yield name, unsharded_tensor
+
+        # TODO: support VeOmni LoRA
+        return param_generator(), None
+
 
 class EngineEvalModeCtx(BaseEngineCtx):
     def __init__(self, engine: VeOmniEngine, **kwargs):
@@ -382,6 +463,8 @@ def __enter__(self):
         assert isinstance(self.engine, VeOmniEngine)
         super().__enter__()
         self.engine.ulysses_sharding_manager.__enter__()
+        # TODO: Switch to eval mode after Integrating the CI environment
+        # VeOmni (ref: https://github.com/ByteDance-Seed/VeOmni/pull/421)
         self.engine.module.train()
 
     def __exit__(self, exc_type, exc_value, traceback):