Offload routed expert of RL

xtuner/v1/data_proto/sequence_context.py

@@ -49,6 +49,7 @@ class SequenceContext:
 
     # moe routed_experts
     rollout_routed_experts: torch.Tensor | None
+    offload_rollout_routed_experts: bool
 
     # Private backing attributes for SP shard reconstruction
     _raw_input_ids: torch.LongTensor | None
@@ -77,6 +78,7 @@ def __init__(
         inputs_embeds: torch.FloatTensor | None = None,
         num_img_tokens: list[list[int]] | None = None,
         rollout_routed_experts: torch.Tensor | None = None,
+        offload_rollout_routed_experts: bool = False,
         # SP shard metadata: private, accessed via properties below
         raw_input_ids: torch.LongTensor | None = None,
         raw_inputs_embeds: torch.FloatTensor | None = None,
@@ -110,6 +112,7 @@ def __init__(
         self.inputs_embeds = inputs_embeds
         self.num_img_tokens = num_img_tokens
         self.rollout_routed_experts = rollout_routed_experts
+        self.offload_rollout_routed_experts = offload_rollout_routed_experts
         self._raw_input_ids = raw_input_ids
         self._raw_inputs_embeds = raw_inputs_embeds
         self._shard_start = shard_start
@@ -232,6 +235,7 @@ def split(self, sequence_parallel_mesh: DeviceMesh | None = None) -> Self:
                 inputs_embeds=self.inputs_embeds,
                 num_img_tokens=self.num_img_tokens,
                 rollout_routed_experts=self.rollout_routed_experts,
+                offload_rollout_routed_experts=self.offload_rollout_routed_experts,
                 raw_input_ids=cast(torch.LongTensor, pad_input_ids),
                 shard_start=start,
                 shard_size=shard_size,
@@ -258,6 +262,7 @@ def cat(cls, sequence_context_list: list["SequenceContext"]) -> "SequenceContext
         num_img_tokens = []
         position_ids = []
         rollout_routed_experts = []
+        offload_rollout_routed_experts = False
 
         for seq_ctx in sequence_context_list:
             assert seq_ctx.sequence_parallel_mesh is None
@@ -287,6 +292,7 @@ def cat(cls, sequence_context_list: list["SequenceContext"]) -> "SequenceContext
                 num_img_tokens.extend(seq_ctx.num_img_tokens)
             if seq_ctx.rollout_routed_experts is not None:
                 rollout_routed_experts.append(seq_ctx.rollout_routed_experts)
+            offload_rollout_routed_experts = offload_rollout_routed_experts or seq_ctx.offload_rollout_routed_experts
             position_ids.append(seq_ctx.position_ids)
         assert len(set(device)) == 1, f"All sequence contexts must be on the same device. Got {set(device)}"
 
@@ -310,6 +316,7 @@ def cat(cls, sequence_context_list: list["SequenceContext"]) -> "SequenceContext
             num_img_tokens=num_img_tokens if num_img_tokens else None,
             position_ids=torch.cat(position_ids, dim=-1) if position_ids else None,  # type: ignore
             rollout_routed_experts=rollout_routed_experts if len(rollout_routed_experts) > 0 else None,  # type: ignore
+            offload_rollout_routed_experts=offload_rollout_routed_experts,
         )
 
     @property
@@ -474,6 +481,9 @@ def copy(self, **overrides) -> Self:
             inputs_embeds=overrides.get("inputs_embeds", self.inputs_embeds),
             num_img_tokens=overrides.get("num_img_tokens", self.num_img_tokens),
             rollout_routed_experts=overrides.get("rollout_routed_experts", self.rollout_routed_experts),
+            offload_rollout_routed_experts=overrides.get(
+                "offload_rollout_routed_experts", self.offload_rollout_routed_experts
+            ),
             raw_input_ids=overrides.get("raw_input_ids", self._raw_input_ids),
             raw_inputs_embeds=overrides.get("raw_inputs_embeds", self._raw_inputs_embeds),
             shard_start=overrides.get("shard_start", self._shard_start),
@@ -528,7 +538,11 @@ def to(self, device: torch.device | str):
         if self.image_grid_thw is not None and hasattr(self.image_grid_thw, "to"):
             self.image_grid_thw = self.image_grid_thw.to(device)  # type: ignore
 
-        if self.rollout_routed_experts is not None and hasattr(self.rollout_routed_experts, "to"):
+        if (
+            self.rollout_routed_experts is not None
+            and not self.offload_rollout_routed_experts
+            and hasattr(self.rollout_routed_experts, "to")
+        ):
             self.rollout_routed_experts = self.rollout_routed_experts.to(device)  # type: ignore
 
         self.device = device
@@ -560,4 +574,5 @@ def data(self) -> dict:
             "inputs_embeds": self.inputs_embeds,
             "num_img_tokens": self.num_img_tokens,
             "rollout_routed_experts": self.rollout_routed_experts,
+            "offload_rollout_routed_experts": self.offload_rollout_routed_experts,
         }

xtuner/v1/module/decoder_layer/moe_decoder_layer.py

@@ -640,6 +640,11 @@ def _pre_moe_forward(
 
         if seq_ctx.rollout_routed_experts is not None and self.layer_idx < seq_ctx.rollout_routed_experts.shape[1]:
             rollout_routed_experts = seq_ctx.rollout_routed_experts[:, self.layer_idx, :]  # seq_l, expert
+            # TODO: pin_memory() + to(device, non_blocking=True) on a CUDA stream would allow overlapping the transfer
+            # with prior-layer compute
+            if seq_ctx.offload_rollout_routed_experts and rollout_routed_experts.device != hidden_states.device:
+                rollout_routed_experts = rollout_routed_experts.contiguous()
+                rollout_routed_experts = rollout_routed_experts.to(hidden_states.device)
         else:
             rollout_routed_experts = None
         router_results: RouterResults = self.gate(hidden_states, rollout_routed_experts)

xtuner/v1/rl/trainer/worker.py

@@ -109,6 +109,8 @@ class WorkerConfig(BaseModel):
             updated weights to rollout workers. Defaults to 0.5.
         seed (int | None): Training worker random seed. When None, the RL
             trainer seed is used. Defaults to None.
+        offload_rollout_routed_experts (bool): Keep rollout routed experts on
+            CPU and move each layer's slice to device during forward. Defaults to False.
 
     **Examples:**
 
@@ -148,6 +150,7 @@ class WorkerConfig(BaseModel):
     profile_time: bool = True
     profile_memory: bool = False
     free_rollout_routed_experts_in_worker: bool = True  # 默认不需要用户配置
+    offload_rollout_routed_experts: bool = False
 
     # sft config
     sft_dataloader_cfg: DataloaderConfig | None = None
@@ -187,6 +190,8 @@ class WorkerTrainLogItem(TypedDict, total=False):
     step_consumed_tokens: int
     efficient_attn_ratio: float
     grad_norm: float
+    max_memory: float
+    reserved_memory: float
 
 
 class WorkerLogItem(TypedDict):
@@ -566,6 +571,7 @@ def fit(self, data_batches: list[WorkerInputItem], rollout_idx: int) -> WorkerLo
             if rollout_routed_experts is not None:
                 self._add_rollout_routed_experts(seq_ctx, rollout_routed_experts)
 
+            seq_ctx.offload_rollout_routed_experts = self.config.offload_rollout_routed_experts
             seq_ctx = data["seq_ctx"].to(DEVICE)
             if self.sp_mesh.size() > 1:
                 seq_ctx = seq_ctx.split(self.sp_mesh)
@@ -788,12 +794,16 @@ def fit(self, data_batches: list[WorkerInputItem], rollout_idx: int) -> WorkerLo
             }
             extra_info_dict = finalize_train_policy_metrics(extra_info_dict, DEVICE)
             train_step_info.pop("total_loss")  # type: ignore[misc]
+            max_memory = DEVICE_MODULE.max_memory_allocated() / (1024**3)  # type: ignore[attr-defined]
+            reserved_memory = DEVICE_MODULE.max_memory_reserved() / (1024**3)  # type: ignore[attr-defined]
 
             train_log_item = WorkerTrainLogItem(
                 **engine_logs_info,  # type: ignore[typeddict-item]
                 **train_step_info,
                 **extra_info_dict,
                 grad_norm=grad_norm.item(),
+                max_memory=max_memory,
+                reserved_memory=reserved_memory,
             )
             worker_log_item["train_metrics"].append(train_log_item)