feat: CP-local loss and configurable CUDA memory profiling

areal/api/cli_args.py

@@ -2341,6 +2341,25 @@ class SessionTracerConfig:
     )
 
 
+@dataclass
+class MemoryProfilerConfig:
+    """CUDA memory snapshot profiling configuration.
+
+    Attributes:
+        profile_steps: Steps at which to record memory snapshots.
+        max_entries: Max entries for torch.cuda.memory._record_memory_history.
+    """
+
+    profile_steps: list[int] = field(
+        default_factory=lambda: [0, 1],
+        metadata={"help": "List of global steps to capture memory snapshots."},
+    )
+    max_entries: int = field(
+        default=100000,
+        metadata={"help": "Max entries for memory history ring buffer."},
+    )
+
+
 @dataclass
 class PerfTracerConfig:
     """Configuration for perf tracer emission."""
@@ -2607,6 +2626,12 @@ class BaseExperimentConfig:
         default=None,
         metadata={"help": "Performance tracer configuration. None means disabled."},
     )
+    memory_profiler: MemoryProfilerConfig | None = field(
+        default=None,
+        metadata={
+            "help": "Memory snapshot profiler configuration. None means disabled."
+        },
+    )
     recover: RecoverConfig = field(default_factory=RecoverConfig)
 
     sglang: SGLangConfig = field(default_factory=SGLangConfig)

areal/api/engine_api.py

@@ -510,6 +510,12 @@ def offload(self) -> None:
     def get_device_stats(self) -> DeviceRuntimeInfo:
         raise NotImplementedError()
 
+    def start_memory_profile(self, max_entries: int = 100000) -> None:
+        pass
+
+    def stop_memory_profile(self, snapshot_dir: str) -> None:
+        pass
+
     def save_perf_tracer(self, step: int | None = None, force: bool = False) -> None:
         """Save performance tracer data.
 

areal/dataset/hhrlhf.py

@@ -21,8 +21,10 @@ def process(sample):
     if max_length is not None:
         # Filter out sequences longer than max_length
         dataset = dataset.filter(
-            lambda x: (len(x["chosen_ids"]) <= max_length)
-            and (len(x["rejected_ids"]) <= max_length)
+            lambda x: (
+                (len(x["chosen_ids"]) <= max_length)
+                and (len(x["rejected_ids"]) <= max_length)
+            )
         )
 
     return dataset

areal/engine/megatron_engine.py

@@ -65,6 +65,7 @@
 from areal.engine.megatron_utils.megatron_lora import get_vllm_lora_target_modules
 from areal.engine.megatron_utils.packed_context_parallel import (
     packed_context_parallel_forward,
+    split_packed_seqs_for_context_parallel,
 )
 from areal.engine.megatron_utils.pipeline_parallel import (
     configure_pipeline_layer_splits,
@@ -713,7 +714,14 @@ def forward_step(batch_iter, model):
                     mb_input.padded_mb.update(tree_kwargs)
                     tree_attn_keys = list(tree_kwargs.keys())
 
-            output = packed_context_parallel_forward(model, mb_input.padded_mb)
+            cp_size = mpu.get_context_parallel_world_size()
+            cp_local = cp_size > 1
+
+            output = packed_context_parallel_forward(
+                model,
+                mb_input.padded_mb,
+                gather_cp_output=not cp_local,
+            )
 
             # Release tree attention metadata after forward pass
             for key in tree_attn_keys:
@@ -729,12 +737,30 @@ def _process_output(input_, output_):
             if mpu.is_pipeline_last_stage(
                 ignore_virtual=False, vp_stage=model_vp_stage
             ):
-                output = unpad_logits(
-                    output,
-                    padding_length=mb_input.padding_length,
-                    cu_seqlens=cu_seqlens,
-                    old_cu_seqlens=mb_input.old_cu_seqlens,
-                )
+                if cp_local and cu_seqlens is not None:
+                    padded_cu_seqlens = mb_input.padded_mb["cu_seqlens"]
+                    rolled_ids = torch.roll(
+                        mb_input.padded_mb["input_ids"], shifts=-1, dims=-1
+                    )
+                    cp_labels = split_packed_seqs_for_context_parallel(
+                        rolled_ids, padded_cu_seqlens
+                    )
+                    cp_loss_mask = split_packed_seqs_for_context_parallel(
+                        mb_input.padded_mb["loss_mask"], padded_cu_seqlens
+                    )
+                    cp_cu_seqlens = padded_cu_seqlens // cp_size
+                    cp_inputs = dict(mb_input.orig_mb)
+                    cp_inputs["_cp_local_labels"] = cp_labels
+                    cp_inputs["loss_mask"] = cp_loss_mask
+                    cp_inputs["cu_seqlens"] = cp_cu_seqlens
+                    return output, functools.partial(_process_output, cp_inputs)
+                else:
+                    output = unpad_logits(
+                        output,
+                        padding_length=mb_input.padding_length,
+                        cu_seqlens=cu_seqlens,
+                        old_cu_seqlens=mb_input.old_cu_seqlens,
+                    )
             return output, functools.partial(_process_output, mb_input.orig_mb)
 
         forward_backward_func = get_forward_backward_func()
@@ -789,7 +815,11 @@ def process_output(
                 loss_multiplier=loss_multiplier,
             )
 
-        self.forward_backward_batch(mb_list, process_output, forward_only=False)
+        self.forward_backward_batch(
+            mb_list,
+            process_output,
+            forward_only=False,
+        )
 
         # Step 4: Optimizer step
         return self.optimizer_step()
@@ -829,7 +859,9 @@ def process_output(
 
         # Step 4: Aggregate losses
         if mpu.is_pipeline_last_stage():
-            return aggregate_eval_losses(losses, mpu.get_data_parallel_group())
+            return aggregate_eval_losses(
+                losses, mpu.get_data_parallel_group(with_context_parallel=True)
+            )
         return None
 
     @torch.no_grad()
@@ -1029,6 +1061,19 @@ def _validate_fp8_consistency(self):
     def get_device_stats(self) -> DeviceRuntimeInfo:
         return DeviceRuntimeInfo.get_current()
 
+    def start_memory_profile(self, max_entries: int = 100000) -> None:
+        torch.cuda.memory._record_memory_history(max_entries=max_entries)
+
+    def stop_memory_profile(self, snapshot_dir: str) -> None:
+        pp = mpu.get_pipeline_model_parallel_rank()
+        dp = mpu.get_data_parallel_rank()
+        cp = mpu.get_context_parallel_rank()
+        tp = mpu.get_tensor_model_parallel_rank()
+        filename = f"snapshot_rank{self.rank:02d}_p{pp}d{dp}c{cp}t{tp}.pickle"
+        path = os.path.join(snapshot_dir, filename)
+        torch.cuda.memory._dump_snapshot(path)
+        torch.cuda.memory._record_memory_history(enabled=None)
+
     def save_perf_tracer(self, step: int | None = None, force: bool = False) -> None:
         perf_tracer.save(step=step, force=force)
 
@@ -1741,7 +1786,11 @@ def _compute_logprobs_and_loss(
                     else None,
                 )
             else:
-                labels = torch.roll(inputs["input_ids"], shifts=-1, dims=-1)
+                cp_local_labels = inputs.get("_cp_local_labels")
+                if cp_local_labels is not None:
+                    labels = cp_local_labels
+                else:
+                    labels = torch.roll(inputs["input_ids"], shifts=-1, dims=-1)
                 logprobs, entropy = gather_logprobs_entropy(
                     output,
                     labels,

areal/engine/megatron_utils/packed_context_parallel.py

@@ -68,10 +68,47 @@ def preprocess_packed_seqs_context_parallel(
     return splitted.unsqueeze(0), packed_seq_params
 
 
+def split_packed_seqs_for_context_parallel(
+    tensor: torch.Tensor,
+    cu_seqlens: torch.Tensor,
+) -> torch.Tensor:
+    """Split a 1D packed tensor using the same interleaved pattern as
+    preprocess_packed_seqs_context_parallel."""
+    cp_size = mpu.get_context_parallel_world_size()
+    cp_rank = mpu.get_context_parallel_rank()
+    if cp_size <= 1:
+        return tensor
+
+    input_lens = cu_seqlens[1:] - cu_seqlens[:-1]
+    batch_size = input_lens.shape[0]
+    output_len = input_lens.sum().item() // cp_size
+
+    splitted = torch.zeros(output_len, dtype=tensor.dtype, device=tensor.device)
+    for i in range(batch_size):
+        seqlen = input_lens[i] // cp_size
+        half_seqlen = seqlen // 2
+        start_idx = cu_seqlens[i] // cp_size
+
+        d = tensor[cu_seqlens[i] : cu_seqlens[i + 1]]
+        splitted[start_idx : start_idx + half_seqlen] = d[
+            half_seqlen * cp_rank : half_seqlen * (cp_rank + 1)
+        ]
+
+        remain_start = input_lens[i] - half_seqlen * (cp_rank + 1)
+        remain_end = input_lens[i] - half_seqlen * cp_rank
+        remain_end = min(remain_end, d.shape[0])
+        remain_len = remain_end - remain_start
+        splitted[start_idx + half_seqlen : start_idx + half_seqlen + remain_len] = d[
+            remain_start:remain_end
+        ]
+    return splitted
+
+
 def postprocess_packed_seqs_context_parallel(
     output: torch.Tensor,
     cu_seqlens: torch.Tensor | None,
     post_process: bool,
+    gather_output: bool = True,
 ) -> torch.Tensor:
     """
     Postprocess packed sequences
@@ -81,6 +118,10 @@ def postprocess_packed_seqs_context_parallel(
         return output
     if cp_size <= 1 or cu_seqlens is None:
         return output.squeeze(0)
+
+    if not gather_output:
+        return output.squeeze(0)
+
     # shape = [batch_size, seq_len] + list(output.shape[2:])
     # [1, packed, dim] -> [batch_size, seq_len, dim]
     batch_size = cu_seqlens.shape[0] - 1
@@ -125,6 +166,7 @@ def postprocess_packed_seqs_context_parallel(
 def packed_context_parallel_forward(
     model: torch.nn.Module,
     input_: dict[str, Any],
+    gather_cp_output: bool = True,
 ):
     input_ids = input_["input_ids"]
     position_ids = input_["position_ids"]
@@ -170,6 +212,6 @@ def packed_context_parallel_forward(
         ignore_virtual=False, vp_stage=model_vp_stage
     )
     output = postprocess_packed_seqs_context_parallel(
-        output, cu_seqlens, is_pipeline_last_stage
+        output, cu_seqlens, is_pipeline_last_stage, gather_output=gather_cp_output
     )
     return output

areal/infra/controller/train_controller.py

@@ -700,6 +700,12 @@ def onload(self) -> None:
     def get_device_stats(self):
         return self._custom_function_call("get_device_stats")
 
+    def start_memory_profile(self, max_entries: int = 100000):
+        return self._custom_function_call("start_memory_profile", max_entries)
+
+    def stop_memory_profile(self, snapshot_dir: str):
+        return self._custom_function_call("stop_memory_profile", snapshot_dir)
+
     def config_perf_tracer(self, config: PerfTracerConfig, role: str) -> None:
         async def _call():
             tasks = [

areal/trainer/rl_trainer.py

@@ -654,6 +654,12 @@ def train(
             # Wait for async checkpoint staging to complete before modifying parameters
             self.saver.maybe_wait_for_staging()
 
+            if (
+                config.memory_profiler is not None
+                and global_step in config.memory_profiler.profile_steps
+            ):
+                self.actor.start_memory_profile(config.memory_profiler.max_entries)
+
             with (
                 stats_tracker.record_timing("train_step"),
                 perf_tracer.trace_scope(
@@ -665,7 +671,18 @@ def train(
                 self.actor.ppo_update(adv_batch)
                 self.actor.step_lr_scheduler()
                 self.actor.get_device_stats().log("ppo update")
-            # Actor stays onloaded through paused window below
+
+            if (
+                config.memory_profiler is not None
+                and global_step in config.memory_profiler.profile_steps
+            ):
+                log_dir = StatsLogger.get_log_path(config.stats_logger)
+                snapshot_dir = os.path.join(
+                    log_dir, "memory_snapshots", f"step_{global_step}"
+                )
+                os.makedirs(snapshot_dir, exist_ok=True)
+                self.actor.stop_memory_profile(snapshot_dir)
+                logger.info(f"Memory snapshots saved to {snapshot_dir}")
 
             if self.critic is not None:
                 with (

areal/trainer/rw_trainer.py

@@ -214,6 +214,12 @@ def train(self):
             # Wait for async checkpoint staging to complete before modifying parameters
             self.saver.maybe_wait_for_staging()
 
+            if (
+                config.memory_profiler is not None
+                and global_step in config.memory_profiler.profile_steps
+            ):
+                self.actor.start_memory_profile(config.memory_profiler.max_entries)
+
             with (
                 stats_tracker.record_timing("train_step"),
                 perf_tracer.trace_scope(
@@ -226,6 +232,18 @@ def train(self):
                 self.actor.step_lr_scheduler()
                 self.actor.get_device_stats().log("after train step")
 
+            if (
+                config.memory_profiler is not None
+                and global_step in config.memory_profiler.profile_steps
+            ):
+                log_dir = StatsLogger.get_log_path(config.stats_logger)
+                snapshot_dir = os.path.join(
+                    log_dir, "memory_snapshots", f"step_{global_step}"
+                )
+                os.makedirs(snapshot_dir, exist_ok=True)
+                self.actor.stop_memory_profile(snapshot_dir)
+                logger.info(f"Memory snapshots saved to {snapshot_dir}")
+
             self.actor.set_version(global_step + 1)
 
             with (

areal/trainer/sft/lm_engine.py

@@ -26,6 +26,7 @@ def train_lm(self, data: list[dict[str, Any]]) -> None:
 
     def _train_lm(self, data: dict[str, Any]) -> None:
         self.engine.train()
+        data["loss_mask"] = torch.roll(data["loss_mask"].bool(), shifts=-1, dims=-1)
         stats = self.engine.train_batch(
             input_=data,
             loss_fn=compute_packed_sft_loss,
@@ -40,6 +41,7 @@ def evaluate_lm(self, data: list[dict[str, Any]]) -> None:
 
     def _evaluate_lm(self, data: dict[str, Any]) -> None:
         self.engine.eval()
+        data["loss_mask"] = torch.roll(data["loss_mask"].bool(), shifts=-1, dims=-1)
         self.engine.eval_batch(
             input_=data,
             loss_fn=compute_packed_sft_loss,
@@ -88,11 +90,10 @@ def compute_packed_sft_loss(
     cu_seqlens: torch.Tensor = input_["cu_seqlens"]
     loss_mask = input_["loss_mask"].bool()
 
-    loss_mask = torch.roll(loss_mask, shifts=-1, dims=-1)
     logprobs = torch.where(loss_mask, logprobs, 0)
 
     device = logprobs.device
-    loss = -logprobs.sum() / loss_mask.count_nonzero()
+    loss = -logprobs.sum() / (1e-5 + loss_mask.count_nonzero())
     with torch.no_grad():
         batch_size = cu_seqlens.shape[0] - 1
         seqlogp = torch.zeros(batch_size, dtype=torch.float64, device=device)

areal/trainer/sft_trainer.py

@@ -180,6 +180,12 @@ def train(self):
 
             self.saver.maybe_wait_for_staging()
 
+            if (
+                config.memory_profiler is not None
+                and global_step in config.memory_profiler.profile_steps
+            ):
+                self.actor.start_memory_profile(config.memory_profiler.max_entries)
+
             with (
                 stats_tracker.record_timing("train_step"),
                 perf_tracer.trace_scope(
@@ -192,6 +198,18 @@ def train(self):
                 self.actor.step_lr_scheduler()
                 self.actor.get_device_stats().log("after train step")
 
+            if (
+                config.memory_profiler is not None
+                and global_step in config.memory_profiler.profile_steps
+            ):
+                log_dir = StatsLogger.get_log_path(config.stats_logger)
+                snapshot_dir = os.path.join(
+                    log_dir, "memory_snapshots", f"step_{global_step}"
+                )
+                os.makedirs(snapshot_dir, exist_ok=True)
+                self.actor.stop_memory_profile(snapshot_dir)
+                logger.info(f"Memory snapshots saved to {snapshot_dir}")
+
             self.actor.set_version(global_step + 1)
 
             with (