feat: CP-local loss and configurable CUDA memory profiling

[yulangz]

Summary

This PR introduces two independent features for Megatron-based SFT training:

1. CP-local cross-entropy loss — eliminates the expensive logits all-gather across Context Parallel (CP) ranks
2. Configurable CUDA memory snapshot profiling — torch.cuda.memory based flame-graph profiling, controllable via YAML config

---

Feature 1: CP-Local Loss

Problem

When Context Parallelism (CP) is enabled, the standard post-processing in postprocess_packed_seqs_context_parallel performs an all-gather to reconstruct the full logits tensor across all CP ranks. For large-vocabulary models (e.g., Qwen3-30B with vocab_size=151,936), this causes OOM because the full [total_tokens, vocab_size] logits tensor must be materialized on each rank.

For a Qwen3-30B-A3B SFT job with seq_len=128K and CP=2:

• Full logits tensor: 128K × 151,936 × bf16 ≈ 37 GiB per rank — exceeds 80 GiB GPU memory

Solution

Compute cross-entropy loss locally on each CP rank using only the local logit shard, avoiding the all-gather entirely.

Key changes:

• packed_context_parallel.py: Added split_packed_seqs_for_context_parallel() to split labels/loss_mask using the same interleaved pattern as the input splitting. Added gather_output flag to postprocess_packed_seqs_context_parallel() to skip all-gather.
• megatron_engine.py:
  • forward_backward_batch() accepts cp_local_loss flag
  • When enabled, labels and loss_mask are pre-rolled and split to match the local CP shard before passing to the loss function
  • Local cu_seqlens are scaled by 1/cp_size to match the local token count
  • _cp_local_labels and _loss_mask_pre_rolled keys are injected into the input dict
  • eval_batch uses dp_group(with_context_parallel=True) for correct loss aggregation across DP×CP ranks

Why this is correct: Each CP rank holds a contiguous interleaved chunk of each sequence. By pre-computing labels (roll(-1) on the full sequence, then split) before the forward pass, each rank can independently compute its local cross-entropy. The global loss is the mean over all valid tokens across all CP ranks, which equals the sum of local losses (weighted by local valid token count) — mathematically equivalent to the all-gather approach.

Files changed

• areal/engine/megatron_utils/packed_context_parallel.py — new split_packed_seqs_for_context_parallel, gather_output parameter
• areal/engine/megatron_engine.py — cp_local_loss logic in forward_backward_batch, train_batch, eval_batch, label routing in _process_megatron_output

---

Feature 2: Configurable Memory Profiling

Problem

CUDA memory snapshot profiling (flame graph) was hardcoded to always run on steps 0 and 1. This should be a configurable feature, disabled by default.

Solution

Follow the existing perf_tracer configuration pattern: add an optional config dataclass, None means disabled.

Config:

@dataclass
class MemoryProfileConfig:
    profile_steps: list[int]  # default [0, 1]
    max_entries: int           # default 100000

Usage in YAML (opt-in):

memory_profile:
  profile_steps: [0, 1]
  max_entries: 100000

No memory_profile key → feature is completely disabled (default).

Key changes:

• areal/api/cli_args.py — new MemoryProfileConfig, added memory_profile field to BaseExperimentConfig
• areal/api/engine_api.py — added start_memory_profile(max_entries) and stop_memory_profile(path) API
• areal/engine/megatron_engine.py — implemented start_memory_profile / stop_memory_profile using torch.cuda.memory._record_memory_history / _dump_snapshot
• areal/trainer/sft/lm_engine.py — RPC forwarding for start_memory_profile(max_entries)
• areal/trainer/sft_trainer.py — config-driven profiling in training loop, snapshots saved to {log_dir}/memory_snapshots/step_{N}.pickle

---

Test

For Qwen3-30B-A3B，In 16K context length and (attn:d1p1t2c4|ffn:d1p1t1e8) parallel, cp local loss can save approximately 5GB memory, and this value is directly proportional to the length of the context.

(with cp local loss) (without cp local loss)

[gemini-code-assist]

Code Review

This pull request introduces CUDA memory snapshot profiling and implements local loss calculation for context parallel (CP) training. Key changes include the addition of a MemoryProfileConfig dataclass, the split_packed_seqs_for_context_parallel utility for handling interleaved packed sequences, and updates to the Megatron engine to support CP-local labels and masks. Feedback suggests making the memory profiling rank configurable instead of hardcoding it to rank 0 to better support distributed debugging. Additionally, it is recommended to vectorize the splitting logic in split_packed_seqs_for_context_parallel to avoid potential performance bottlenecks caused by Python loops during microbatch processing.

[gemini-code-assist]

Hardcoding self.rank == 0 for memory profiling is restrictive. In distributed training with Megatron (PP/TP/CP), memory pressure and OOM issues often vary significantly across different ranks (e.g., the last pipeline stage or specific tensor parallel ranks). Consider making the profiling rank configurable or allowing it on the local rank 0 of each node to provide better visibility into the cluster's memory state.

[gemini-code-assist]

The CP-local loss logic performs torch.roll and split_packed_seqs_for_context_parallel (which contains a Python loop over the batch size) for every microbatch in the last pipeline stage. While functional, this could become a performance bottleneck for large batch sizes or high-frequency microbatching. Consider vectorizing the interleaved splitting logic in split_packed_seqs_for_context_parallel to avoid the Python loop.

[yulangz]

Agreed that vectorization is possible here, but split_packed_seqs_for_context_parallel is not a hot path — it runs once per micro-batch during loss computation, not in the inner attention loop. The current loop-based implementation is clear and correct. We can optimize this in a follow-up if profiling shows it matters.

[Copilot]

Pull request overview

This PR adds two opt-in improvements for Megatron-based SFT training: a context-parallel (CP) local loss path to avoid CP logits all-gather, and a YAML-configurable CUDA memory snapshot profiler.

Changes:

• Introduces CP-local packed-sequence splitting utilities and an option to skip CP all-gather in packed forward/postprocess.
• Updates Megatron engine training/eval flow to support CP-local label/mask routing and correct loss aggregation in eval.
• Adds a memory_profile config and wiring to start/stop CUDA memory snapshot recording from the SFT training loop.

Reviewed changes

Copilot reviewed 7 out of 7 changed files in this pull request and generated 5 comments.

Show a summary per file

File	Description
areal/trainer/sft_trainer.py	Triggers memory profiling start/stop and snapshot writing per configured steps.
areal/trainer/sft/lm_engine.py	Adds controller RPC hooks for memory profiling; adjusts SFT loss mask rolling behavior.
areal/engine/megatron_utils/packed_context_parallel.py	Adds packed-tensor CP split helper and a flag to skip CP output all-gather.
areal/engine/megatron_engine.py	Implements CP-local loss routing and adds engine-level memory snapshot APIs.
areal/api/engine_api.py	Extends the TrainEngine API surface with memory profiling hooks.
areal/api/cli_args.py	Adds MemoryProfileConfig and a memory_profile field to experiment config.
areal/dataset/hhrlhf.py	Formatting-only change to dataset length filtering lambda.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[Copilot]

cp_local_loss is force-enabled whenever context_parallel_world_size() > 1. However, the CP-local path only rolls/splits input_ids (labels) and loss_mask/cu_seqlens, and it also mutates semantics by rolling loss_mask. This will break non-SFT training paths (e.g., PPO/RW) whose loss functions expect additional per-token tensors in inputs (advantages/old_logprobs/returns/etc.) and an un-rolled loss_mask. Make CP-local loss opt-in (plumb a config/flag down to train_batch/eval_batch) or generalize the CP-local input rewriting to split all token-aligned tensors without changing their semantics for other algorithms.

[yulangz]

This is not a concern — train_batch and eval_batch (which call forward_backward_batch) are only used by the SFT path. The RL training path uses different methods (forward_backward_func directly via Megatron pipeline). So the CP-local logic is scoped correctly to SFT only.

In the latest commit, cp_local_loss has been removed as a parameter entirely — the engine now infers CP-local mode internally from mpu.get_context_parallel_world_size() > 1.

[Copilot]

In the CP-local branch, cp_inputs = dict(mb_input.orig_mb) keeps the unsplit original microbatch tensors (and possibly unpadded lengths) but only replaces loss_mask/cu_seqlens and injects _cp_local_labels. Any other packed, token-aligned fields in orig_mb (e.g., PPO/RW training targets) will have mismatched shapes relative to the CP-sharded output, leading to incorrect loss computation or runtime shape errors. Consider building cp_inputs from padded_mb and splitting every token-aligned tensor with split_packed_seqs_for_context_parallel, or restrict this path to the SFT loss that only consumes labels/loss_mask.

[yulangz]

This is by design. cp_inputs = dict(mb_input.orig_mb) creates a shallow copy of the original microbatch dict, then we overwrite the fields that need CP-local values (_cp_local_labels, loss_mask, cu_seqlens). The SFT loss function (compute_packed_sft_loss) only consumes these replaced fields — it does not use any of the original unsplit tensors. The remaining fields in orig_mb (like metadata) are harmlessly carried through.

[Copilot]

SFTTrainer unconditionally calls self.actor.start_memory_profile(...) / stop_memory_profile(...) when the config is set, but the v1 controller (LMController) currently does not define these methods (only LMControllerV2 does). This will raise AttributeError for non-v2 runs. Either add these methods to LMController as well (dispatch via _custom_function_call with rpc_meta={"broadcast": True}), or guard these calls based on controller version/capability.

[Copilot]

loss_mask.count_nonzero() returns a 0-d tensor, so if n_valid == 0: will raise RuntimeError: Boolean value of Tensor... at runtime. Convert to a Python number (e.g., n_valid = int(... .item())) or use n_valid.item() == 0 before the conditional, and keep the tensor form only for arithmetic/division as needed.

[yulangz]

Fixed in the latest commit. The n_valid == 0 branch has been removed entirely — the loss is now computed as loss = -logprobs.sum() / (1e-5 + loss_mask.count_nonzero()), which handles the all-zero case gracefully via the epsilon term.

[Copilot]

LMControllerV2 inherits from GatewayTrainController, which does not define _custom_function_call. These methods will raise AttributeError when invoked (e.g., in SFTTrainer.train). Implement these via gateway endpoints (similar to train_batch/step_lr_scheduler), or route through an existing dispatch helper on GatewayTrainController if one exists.

Suggested change

	self._custom_function_call("start_memory_profile", max_entries)
	def stop_memory_profile(self, path: str):
	self._custom_function_call("stop_memory_profile", path)
	payload = {
	"args": serialize_value([max_entries]),
	"kwargs": serialize_value({}),
	}
	self._gateway_post_result("/sft/start_memory_profile", payload)
	def stop_memory_profile(self, path: str):
	payload = {
	"args": serialize_value([path]),
	"kwargs": serialize_value({}),
	}
	self._gateway_post_result("/sft/stop_memory_profile", payload)

[yulangz]

LMControllerV2.stop_memory_profile and start_memory_profile use _custom_function_call which was already present in the pre-existing code (see the same class's existing methods). This is not introduced by this PR — it's inherited from GatewayTrainController base class.

[garrett4wade]

(1) This method call won't work for the single controller mode because the train controller does not implement this method.
(2) Not implemented for the RL/RW trainer. Should either amend the code or open an issue about this.

[yulangz]

Both issues addressed:

1. Added start_memory_profile and stop_memory_profile to LMController (V1) for single-controller mode.
2. Ported memory profiling support to PPOTrainer (rl_trainer.py) and RWTrainer (rw_trainer.py) with the same config-driven pattern as SFTTrainer.

[garrett4wade]

(1) mpu.get_data_parallel_group(with_context_parallel=True) implies the else branch and is sufficient.
(2) Why should we also gather across the CP group? The data should be duplicated across CP groups. Summing the losses up would upscale the loss. Conversely, is it because we have incorrectly downscaled the training loss?

[yulangz]

Good question. Using with_context_parallel=True is always safe:

1. CP=1: The CP group is trivially size-1, so get_data_parallel_group(with_context_parallel=True) returns the same group as get_data_parallel_group() — it is idempotent.
2. CP>1 (CP-local loss): Each CP rank computes loss over its local shard. aggregate_eval_losses does a SUM all-reduce. With with_context_parallel=True, the reduction spans DP×CP ranks, correctly averaging the CP-local losses across all CP ranks.

On the training side, DistributedOptimizer uses intra_dp_cp_group (DP×CP) for gradient reduce-scatter, so CP-local gradients are mathematically equivalent to the global gradient.

Since with_context_parallel=True is always correct regardless of CP size, we can simplify to always use it. Updated in the latest commit.