From 430a1f9cc1697432812f3482d16e89c0a4dacdfd Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E5=8D=9A=E6=83=9F?= Date: Thu, 23 Apr 2026 16:29:41 +0800 Subject: [PATCH] feat(engine): add built-in memory profiler support and avoid gathering CP logits --- areal/api/cli_args.py | 25 +++ areal/api/engine_api.py | 6 + areal/dataset/hhrlhf.py | 6 +- areal/engine/megatron_engine.py | 69 ++++++-- .../megatron_utils/packed_context_parallel.py | 44 ++++- areal/infra/controller/train_controller.py | 6 + areal/trainer/rl_trainer.py | 19 ++- areal/trainer/rw_trainer.py | 18 ++ areal/trainer/sft/lm_engine.py | 5 +- areal/trainer/sft_trainer.py | 18 ++ docs/en/cli_reference.md | 156 ++++++++++-------- docs/zh/cli_reference.md | 156 ++++++++++-------- 12 files changed, 376 insertions(+), 152 deletions(-) diff --git a/areal/api/cli_args.py b/areal/api/cli_args.py index e05998934f..7e08f2d63a 100644 --- a/areal/api/cli_args.py +++ b/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) diff --git a/areal/api/engine_api.py b/areal/api/engine_api.py index 2e4b0ec098..08cd31c45c 100644 --- a/areal/api/engine_api.py +++ b/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. diff --git a/areal/dataset/hhrlhf.py b/areal/dataset/hhrlhf.py index fcb150ad1e..12d08346ab 100644 --- a/areal/dataset/hhrlhf.py +++ b/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 diff --git a/areal/engine/megatron_engine.py b/areal/engine/megatron_engine.py index 53737d5499..fd093f79fd 100644 --- a/areal/engine/megatron_engine.py +++ b/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, diff --git a/areal/engine/megatron_utils/packed_context_parallel.py b/areal/engine/megatron_utils/packed_context_parallel.py index 3eeb1c770b..1063a91295 100644 --- a/areal/engine/megatron_utils/packed_context_parallel.py +++ b/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 diff --git a/areal/infra/controller/train_controller.py b/areal/infra/controller/train_controller.py index a7c493d5a8..4ee4e1854e 100644 --- a/areal/infra/controller/train_controller.py +++ b/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 = [ diff --git a/areal/trainer/rl_trainer.py b/areal/trainer/rl_trainer.py index 3bed28f443..978a6c18de 100644 --- a/areal/trainer/rl_trainer.py +++ b/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 ( diff --git a/areal/trainer/rw_trainer.py b/areal/trainer/rw_trainer.py index efabb7263f..51bd239150 100644 --- a/areal/trainer/rw_trainer.py +++ b/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 ( diff --git a/areal/trainer/sft/lm_engine.py b/areal/trainer/sft/lm_engine.py index 4551a4ae7b..9dd3f8439e 100644 --- a/areal/trainer/sft/lm_engine.py +++ b/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) diff --git a/areal/trainer/sft_trainer.py b/areal/trainer/sft_trainer.py index c7b1f52275..27b1598546 100644 --- a/areal/trainer/sft_trainer.py +++ b/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 ( diff --git a/docs/en/cli_reference.md b/docs/en/cli_reference.md index 0fbd2f8098..b567657fa6 100644 --- a/docs/en/cli_reference.md +++ b/docs/en/cli_reference.md @@ -76,6 +76,7 @@ For detailed examples, see the experiment configurations in the `examples/` dire - [DistributedDataParallel Configuration](section-distributed-data-parallel) - [FP8Engine Configuration](section-fp8-engine) - [MegatronEngine Configuration](section-megatron-engine) +- [MemoryProfiler Configuration](section-memory-profiler) - [OpenAIProxy Configuration](section-open-ai-proxy) - [PerfTracer Configuration](section-perf-tracer) - [RejectionSampling Configuration](section-rejection-sampling) @@ -93,28 +94,29 @@ ______________________________________________________________________ Base configuration class for all experiment types with common settings. -| Parameter | Type | Default | Description | -| -------------------- | ----------------------------------------------------- | ------------ | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | -| `experiment_name` | string | **Required** | Name of the experiment (no '\_' or '/'). Required. | -| `trial_name` | string | **Required** | Name of the trial (no '-' or '/'). Required. | -| `cluster` | [`ClusterSpecConfig`](section-cluster) | **Required** | Cluster specification. Mainly used by slurm. | -| `allocation_mode` | string | `""` | DEPRECATED: Use per-engine 'backend' fields instead (e.g., actor.backend, rollout.backend). Legacy pattern-based GPU parallel strategy allocation mode. Only used by SPMD launchers (local/ray/slurm). Manual migration to per-engine 'backend' fields is required. | -| `seed` | integer | `1` | Random seed for reproducibility. | -| `enable_offload` | boolean | `False` | Whether to enable training offload using torch_memory_saver. This requires setting up the environment for TMS (e.g., via LD_PRELOAD). | -| `total_train_epochs` | integer | `1` | Total number of epochs to train the model. | -| `total_train_steps` | integer \| None | `None` | Terminate training after this number of steps. For benchmarking purposes only. None indicates normal training. | -| `total_train_n_seqs` | integer \| None | `None` | Terminate training after consuming this number of samples. For benchmarking purposes only. None indicates normal training. | -| `tokenizer_path` | string | `""` | Path to the tokenizer. | -| `train_dataset` | [`TrainDatasetConfig`](section-train-dataset) | **Required** | - | -| `valid_dataset` | [`ValidDatasetConfig`](section-valid-dataset) \| None | `None` | - | -| `saver` | [`SaverConfig`](section-saver) | **Required** | - | -| `evaluator` | [`EvaluatorConfig`](section-evaluator) | **Required** | - | -| `stats_logger` | [`StatsLoggerConfig`](section-stats-logger) | **Required** | - | -| `perf_tracer` | [`PerfTracerConfig`](section-perf-tracer) \| None | `None` | Performance tracer configuration. None means disabled. | -| `recover` | [`RecoverConfig`](section-recover) | **Required** | - | -| `sglang` | [`SGLangConfig`](section-sg-lang) | **Required** | - | -| `vllm` | [`vLLMConfig`](section-v-llm) | **Required** | - | -| `scheduler` | [`SchedulerConfig`](section-scheduler) | **Required** | - | +| Parameter | Type | Default | Description | +| -------------------- | --------------------------------------------------------- | ------------ | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | +| `experiment_name` | string | **Required** | Name of the experiment (no '\_' or '/'). Required. | +| `trial_name` | string | **Required** | Name of the trial (no '-' or '/'). Required. | +| `cluster` | [`ClusterSpecConfig`](section-cluster) | **Required** | Cluster specification. Mainly used by slurm. | +| `allocation_mode` | string | `""` | DEPRECATED: Use per-engine 'backend' fields instead (e.g., actor.backend, rollout.backend). Legacy pattern-based GPU parallel strategy allocation mode. Only used by SPMD launchers (local/ray/slurm). Manual migration to per-engine 'backend' fields is required. | +| `seed` | integer | `1` | Random seed for reproducibility. | +| `enable_offload` | boolean | `False` | Whether to enable training offload using torch_memory_saver. This requires setting up the environment for TMS (e.g., via LD_PRELOAD). | +| `total_train_epochs` | integer | `1` | Total number of epochs to train the model. | +| `total_train_steps` | integer \| None | `None` | Terminate training after this number of steps. For benchmarking purposes only. None indicates normal training. | +| `total_train_n_seqs` | integer \| None | `None` | Terminate training after consuming this number of samples. For benchmarking purposes only. None indicates normal training. | +| `tokenizer_path` | string | `""` | Path to the tokenizer. | +| `train_dataset` | [`TrainDatasetConfig`](section-train-dataset) | **Required** | - | +| `valid_dataset` | [`ValidDatasetConfig`](section-valid-dataset) \| None | `None` | - | +| `saver` | [`SaverConfig`](section-saver) | **Required** | - | +| `evaluator` | [`EvaluatorConfig`](section-evaluator) | **Required** | - | +| `stats_logger` | [`StatsLoggerConfig`](section-stats-logger) | **Required** | - | +| `perf_tracer` | [`PerfTracerConfig`](section-perf-tracer) \| None | `None` | Performance tracer configuration. None means disabled. | +| `memory_profiler` | [`MemoryProfilerConfig`](section-memory-profiler) \| None | `None` | Memory snapshot profiler configuration. None means disabled. | +| `recover` | [`RecoverConfig`](section-recover) | **Required** | - | +| `sglang` | [`SGLangConfig`](section-sg-lang) | **Required** | - | +| `vllm` | [`vLLMConfig`](section-v-llm) | **Required** | - | +| `scheduler` | [`SchedulerConfig`](section-scheduler) | **Required** | - | (section-grpo)= @@ -140,6 +142,7 @@ A dummy place holder of GRPO config for backward compatibility. | `evaluator` | [`EvaluatorConfig`](section-evaluator) | **Required** | - | | `stats_logger` | [`StatsLoggerConfig`](section-stats-logger) | **Required** | - | | `perf_tracer` | [`PerfTracerConfig`](section-perf-tracer) \| None | `None` | Performance tracer configuration. None means disabled. | +| `memory_profiler` | [`MemoryProfilerConfig`](section-memory-profiler) \| None | `None` | Memory snapshot profiler configuration. None means disabled. | | `recover` | [`RecoverConfig`](section-recover) | **Required** | - | | `sglang` | [`SGLangConfig`](section-sg-lang) | **Required** | - | | `vllm` | [`vLLMConfig`](section-v-llm) | **Required** | - | @@ -177,6 +180,7 @@ Configuration for Proximal Policy Optimization (PPO) reinforcement learning expe | `evaluator` | [`EvaluatorConfig`](section-evaluator) | **Required** | - | | `stats_logger` | [`StatsLoggerConfig`](section-stats-logger) | **Required** | - | | `perf_tracer` | [`PerfTracerConfig`](section-perf-tracer) \| None | `None` | Performance tracer configuration. None means disabled. | +| `memory_profiler` | [`MemoryProfilerConfig`](section-memory-profiler) \| None | `None` | Memory snapshot profiler configuration. None means disabled. | | `recover` | [`RecoverConfig`](section-recover) | **Required** | - | | `sglang` | [`SGLangConfig`](section-sg-lang) | **Required** | - | | `vllm` | [`vLLMConfig`](section-v-llm) | **Required** | - | @@ -196,29 +200,30 @@ Configuration for Proximal Policy Optimization (PPO) reinforcement learning expe Configuration for Reward Model (RW) training experiments. -| Parameter | Type | Default | Description | -| -------------------- | ----------------------------------------------------- | ------------ | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | -| `experiment_name` | string | **Required** | Name of the experiment (no '\_' or '/'). Required. | -| `trial_name` | string | **Required** | Name of the trial (no '-' or '/'). Required. | -| `cluster` | [`ClusterSpecConfig`](section-cluster) | **Required** | Cluster specification. Mainly used by slurm. | -| `allocation_mode` | string | `""` | DEPRECATED: Use per-engine 'backend' fields instead (e.g., actor.backend, rollout.backend). Legacy pattern-based GPU parallel strategy allocation mode. Only used by SPMD launchers (local/ray/slurm). Manual migration to per-engine 'backend' fields is required. | -| `seed` | integer | `1` | Random seed for reproducibility. | -| `enable_offload` | boolean | `False` | Whether to enable training offload using torch_memory_saver. This requires setting up the environment for TMS (e.g., via LD_PRELOAD). | -| `total_train_epochs` | integer | `1` | Total number of epochs to train the model. | -| `total_train_steps` | integer \| None | `None` | Terminate training after this number of steps. For benchmarking purposes only. None indicates normal training. | -| `total_train_n_seqs` | integer \| None | `None` | Terminate training after consuming this number of samples. For benchmarking purposes only. None indicates normal training. | -| `tokenizer_path` | string | `""` | Path to the tokenizer. | -| `train_dataset` | [`TrainDatasetConfig`](section-train-dataset) | **Required** | - | -| `valid_dataset` | [`ValidDatasetConfig`](section-valid-dataset) \| None | `None` | - | -| `saver` | [`SaverConfig`](section-saver) | **Required** | - | -| `evaluator` | [`EvaluatorConfig`](section-evaluator) | **Required** | - | -| `stats_logger` | [`StatsLoggerConfig`](section-stats-logger) | **Required** | - | -| `perf_tracer` | [`PerfTracerConfig`](section-perf-tracer) \| None | `None` | Performance tracer configuration. None means disabled. | -| `recover` | [`RecoverConfig`](section-recover) | **Required** | - | -| `sglang` | [`SGLangConfig`](section-sg-lang) | **Required** | - | -| `vllm` | [`vLLMConfig`](section-v-llm) | **Required** | - | -| `scheduler` | [`SchedulerConfig`](section-scheduler) | **Required** | - | -| `actor` | [`TrainEngineConfig`](section-train-engine) | **Required** | - | +| Parameter | Type | Default | Description | +| -------------------- | --------------------------------------------------------- | ------------ | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | +| `experiment_name` | string | **Required** | Name of the experiment (no '\_' or '/'). Required. | +| `trial_name` | string | **Required** | Name of the trial (no '-' or '/'). Required. | +| `cluster` | [`ClusterSpecConfig`](section-cluster) | **Required** | Cluster specification. Mainly used by slurm. | +| `allocation_mode` | string | `""` | DEPRECATED: Use per-engine 'backend' fields instead (e.g., actor.backend, rollout.backend). Legacy pattern-based GPU parallel strategy allocation mode. Only used by SPMD launchers (local/ray/slurm). Manual migration to per-engine 'backend' fields is required. | +| `seed` | integer | `1` | Random seed for reproducibility. | +| `enable_offload` | boolean | `False` | Whether to enable training offload using torch_memory_saver. This requires setting up the environment for TMS (e.g., via LD_PRELOAD). | +| `total_train_epochs` | integer | `1` | Total number of epochs to train the model. | +| `total_train_steps` | integer \| None | `None` | Terminate training after this number of steps. For benchmarking purposes only. None indicates normal training. | +| `total_train_n_seqs` | integer \| None | `None` | Terminate training after consuming this number of samples. For benchmarking purposes only. None indicates normal training. | +| `tokenizer_path` | string | `""` | Path to the tokenizer. | +| `train_dataset` | [`TrainDatasetConfig`](section-train-dataset) | **Required** | - | +| `valid_dataset` | [`ValidDatasetConfig`](section-valid-dataset) \| None | `None` | - | +| `saver` | [`SaverConfig`](section-saver) | **Required** | - | +| `evaluator` | [`EvaluatorConfig`](section-evaluator) | **Required** | - | +| `stats_logger` | [`StatsLoggerConfig`](section-stats-logger) | **Required** | - | +| `perf_tracer` | [`PerfTracerConfig`](section-perf-tracer) \| None | `None` | Performance tracer configuration. None means disabled. | +| `memory_profiler` | [`MemoryProfilerConfig`](section-memory-profiler) \| None | `None` | Memory snapshot profiler configuration. None means disabled. | +| `recover` | [`RecoverConfig`](section-recover) | **Required** | - | +| `sglang` | [`SGLangConfig`](section-sg-lang) | **Required** | - | +| `vllm` | [`vLLMConfig`](section-v-llm) | **Required** | - | +| `scheduler` | [`SchedulerConfig`](section-scheduler) | **Required** | - | +| `actor` | [`TrainEngineConfig`](section-train-engine) | **Required** | - | (section-sft)= @@ -226,29 +231,30 @@ Configuration for Reward Model (RW) training experiments. Configuration for Supervised Fine-Tuning (SFT) experiments. -| Parameter | Type | Default | Description | -| -------------------- | ----------------------------------------------------- | ------------ | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | -| `experiment_name` | string | **Required** | Name of the experiment (no '\_' or '/'). Required. | -| `trial_name` | string | **Required** | Name of the trial (no '-' or '/'). Required. | -| `cluster` | [`ClusterSpecConfig`](section-cluster) | **Required** | Cluster specification. Mainly used by slurm. | -| `allocation_mode` | string | `""` | DEPRECATED: Use per-engine 'backend' fields instead (e.g., actor.backend, rollout.backend). Legacy pattern-based GPU parallel strategy allocation mode. Only used by SPMD launchers (local/ray/slurm). Manual migration to per-engine 'backend' fields is required. | -| `seed` | integer | `1` | Random seed for reproducibility. | -| `enable_offload` | boolean | `False` | Whether to enable training offload using torch_memory_saver. This requires setting up the environment for TMS (e.g., via LD_PRELOAD). | -| `total_train_epochs` | integer | `1` | Total number of epochs to train the model. | -| `total_train_steps` | integer \| None | `None` | Terminate training after this number of steps. For benchmarking purposes only. None indicates normal training. | -| `total_train_n_seqs` | integer \| None | `None` | Terminate training after consuming this number of samples. For benchmarking purposes only. None indicates normal training. | -| `tokenizer_path` | string | `""` | Path to the tokenizer. | -| `train_dataset` | [`TrainDatasetConfig`](section-train-dataset) | **Required** | - | -| `valid_dataset` | [`ValidDatasetConfig`](section-valid-dataset) \| None | `None` | - | -| `saver` | [`SaverConfig`](section-saver) | **Required** | - | -| `evaluator` | [`EvaluatorConfig`](section-evaluator) | **Required** | - | -| `stats_logger` | [`StatsLoggerConfig`](section-stats-logger) | **Required** | - | -| `perf_tracer` | [`PerfTracerConfig`](section-perf-tracer) \| None | `None` | Performance tracer configuration. None means disabled. | -| `recover` | [`RecoverConfig`](section-recover) | **Required** | - | -| `sglang` | [`SGLangConfig`](section-sg-lang) | **Required** | - | -| `vllm` | [`vLLMConfig`](section-v-llm) | **Required** | - | -| `scheduler` | [`SchedulerConfig`](section-scheduler) | **Required** | - | -| `actor` | [`TrainEngineConfig`](section-train-engine) | **Required** | - | +| Parameter | Type | Default | Description | +| -------------------- | --------------------------------------------------------- | ------------ | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | +| `experiment_name` | string | **Required** | Name of the experiment (no '\_' or '/'). Required. | +| `trial_name` | string | **Required** | Name of the trial (no '-' or '/'). Required. | +| `cluster` | [`ClusterSpecConfig`](section-cluster) | **Required** | Cluster specification. Mainly used by slurm. | +| `allocation_mode` | string | `""` | DEPRECATED: Use per-engine 'backend' fields instead (e.g., actor.backend, rollout.backend). Legacy pattern-based GPU parallel strategy allocation mode. Only used by SPMD launchers (local/ray/slurm). Manual migration to per-engine 'backend' fields is required. | +| `seed` | integer | `1` | Random seed for reproducibility. | +| `enable_offload` | boolean | `False` | Whether to enable training offload using torch_memory_saver. This requires setting up the environment for TMS (e.g., via LD_PRELOAD). | +| `total_train_epochs` | integer | `1` | Total number of epochs to train the model. | +| `total_train_steps` | integer \| None | `None` | Terminate training after this number of steps. For benchmarking purposes only. None indicates normal training. | +| `total_train_n_seqs` | integer \| None | `None` | Terminate training after consuming this number of samples. For benchmarking purposes only. None indicates normal training. | +| `tokenizer_path` | string | `""` | Path to the tokenizer. | +| `train_dataset` | [`TrainDatasetConfig`](section-train-dataset) | **Required** | - | +| `valid_dataset` | [`ValidDatasetConfig`](section-valid-dataset) \| None | `None` | - | +| `saver` | [`SaverConfig`](section-saver) | **Required** | - | +| `evaluator` | [`EvaluatorConfig`](section-evaluator) | **Required** | - | +| `stats_logger` | [`StatsLoggerConfig`](section-stats-logger) | **Required** | - | +| `perf_tracer` | [`PerfTracerConfig`](section-perf-tracer) \| None | `None` | Performance tracer configuration. None means disabled. | +| `memory_profiler` | [`MemoryProfilerConfig`](section-memory-profiler) \| None | `None` | Memory snapshot profiler configuration. None means disabled. | +| `recover` | [`RecoverConfig`](section-recover) | **Required** | - | +| `sglang` | [`SGLangConfig`](section-sg-lang) | **Required** | - | +| `vllm` | [`vLLMConfig`](section-v-llm) | **Required** | - | +| `scheduler` | [`SchedulerConfig`](section-scheduler) | **Required** | - | +| `actor` | [`TrainEngineConfig`](section-train-engine) | **Required** | - | (section-fsdp-engine)= @@ -970,6 +976,20 @@ Refer to Megatron-LM documentation for implementation details. | `fp8_config` | [`FP8EngineConfig`](section-fp8-engine) \| None | `None` | - | | `bridge_type` | string | `"mbridge"` | Bridge backend for MegatronEngine. Choices: 'mbridge' or 'megatron-bridge'. **Choices:** `mbridge`, `megatron-bridge` | +(section-memory-profiler)= + +## MemoryProfiler Configuration + +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. + +| Parameter | Type | Default | Description | +| --------------- | --------------- | ------------ | ------------------------------------------------- | +| `profile_steps` | list of integer | **Required** | List of global steps to capture memory snapshots. | +| `max_entries` | integer | `100000` | Max entries for memory history ring buffer. | + (section-open-ai-proxy)= ## OpenAIProxy Configuration diff --git a/docs/zh/cli_reference.md b/docs/zh/cli_reference.md index 17660f016f..9b26f18a9b 100644 --- a/docs/zh/cli_reference.md +++ b/docs/zh/cli_reference.md @@ -74,6 +74,7 @@ python3 train.py --config path/to/config.yaml actor.lr=1e-4 seed=42 - [DistributedDataParallel Configuration](section-distributed-data-parallel) - [FP8Engine Configuration](section-fp8-engine) - [MegatronEngine Configuration](section-megatron-engine) +- [MemoryProfiler Configuration](section-memory-profiler) - [OpenAIProxy Configuration](section-open-ai-proxy) - [PerfTracer Configuration](section-perf-tracer) - [RejectionSampling Configuration](section-rejection-sampling) @@ -91,28 +92,29 @@ ______________________________________________________________________ Base configuration class for all experiment types with common settings. -| Parameter | Type | Default | Description | -| -------------------- | ----------------------------------------------------- | ------------ | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | -| `experiment_name` | string | **Required** | Name of the experiment (no '\_' or '/'). Required. | -| `trial_name` | string | **Required** | Name of the trial (no '-' or '/'). Required. | -| `cluster` | [`ClusterSpecConfig`](section-cluster) | **Required** | Cluster specification. Mainly used by slurm. | -| `allocation_mode` | string | `""` | DEPRECATED: Use per-engine 'backend' fields instead (e.g., actor.backend, rollout.backend). Legacy pattern-based GPU parallel strategy allocation mode. Only used by SPMD launchers (local/ray/slurm). Manual migration to per-engine 'backend' fields is required. | -| `seed` | integer | `1` | Random seed for reproducibility. | -| `enable_offload` | boolean | `False` | Whether to enable training offload using torch_memory_saver. This requires setting up the environment for TMS (e.g., via LD_PRELOAD). | -| `total_train_epochs` | integer | `1` | Total number of epochs to train the model. | -| `total_train_steps` | integer \| None | `None` | Terminate training after this number of steps. For benchmarking purposes only. None indicates normal training. | -| `total_train_n_seqs` | integer \| None | `None` | Terminate training after consuming this number of samples. For benchmarking purposes only. None indicates normal training. | -| `tokenizer_path` | string | `""` | Path to the tokenizer. | -| `train_dataset` | [`TrainDatasetConfig`](section-train-dataset) | **Required** | - | -| `valid_dataset` | [`ValidDatasetConfig`](section-valid-dataset) \| None | `None` | - | -| `saver` | [`SaverConfig`](section-saver) | **Required** | - | -| `evaluator` | [`EvaluatorConfig`](section-evaluator) | **Required** | - | -| `stats_logger` | [`StatsLoggerConfig`](section-stats-logger) | **Required** | - | -| `perf_tracer` | [`PerfTracerConfig`](section-perf-tracer) \| None | `None` | Performance tracer configuration. None means disabled. | -| `recover` | [`RecoverConfig`](section-recover) | **Required** | - | -| `sglang` | [`SGLangConfig`](section-sg-lang) | **Required** | - | -| `vllm` | [`vLLMConfig`](section-v-llm) | **Required** | - | -| `scheduler` | [`SchedulerConfig`](section-scheduler) | **Required** | - | +| Parameter | Type | Default | Description | +| -------------------- | --------------------------------------------------------- | ------------ | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | +| `experiment_name` | string | **Required** | Name of the experiment (no '\_' or '/'). Required. | +| `trial_name` | string | **Required** | Name of the trial (no '-' or '/'). Required. | +| `cluster` | [`ClusterSpecConfig`](section-cluster) | **Required** | Cluster specification. Mainly used by slurm. | +| `allocation_mode` | string | `""` | DEPRECATED: Use per-engine 'backend' fields instead (e.g., actor.backend, rollout.backend). Legacy pattern-based GPU parallel strategy allocation mode. Only used by SPMD launchers (local/ray/slurm). Manual migration to per-engine 'backend' fields is required. | +| `seed` | integer | `1` | Random seed for reproducibility. | +| `enable_offload` | boolean | `False` | Whether to enable training offload using torch_memory_saver. This requires setting up the environment for TMS (e.g., via LD_PRELOAD). | +| `total_train_epochs` | integer | `1` | Total number of epochs to train the model. | +| `total_train_steps` | integer \| None | `None` | Terminate training after this number of steps. For benchmarking purposes only. None indicates normal training. | +| `total_train_n_seqs` | integer \| None | `None` | Terminate training after consuming this number of samples. For benchmarking purposes only. None indicates normal training. | +| `tokenizer_path` | string | `""` | Path to the tokenizer. | +| `train_dataset` | [`TrainDatasetConfig`](section-train-dataset) | **Required** | - | +| `valid_dataset` | [`ValidDatasetConfig`](section-valid-dataset) \| None | `None` | - | +| `saver` | [`SaverConfig`](section-saver) | **Required** | - | +| `evaluator` | [`EvaluatorConfig`](section-evaluator) | **Required** | - | +| `stats_logger` | [`StatsLoggerConfig`](section-stats-logger) | **Required** | - | +| `perf_tracer` | [`PerfTracerConfig`](section-perf-tracer) \| None | `None` | Performance tracer configuration. None means disabled. | +| `memory_profiler` | [`MemoryProfilerConfig`](section-memory-profiler) \| None | `None` | Memory snapshot profiler configuration. None means disabled. | +| `recover` | [`RecoverConfig`](section-recover) | **Required** | - | +| `sglang` | [`SGLangConfig`](section-sg-lang) | **Required** | - | +| `vllm` | [`vLLMConfig`](section-v-llm) | **Required** | - | +| `scheduler` | [`SchedulerConfig`](section-scheduler) | **Required** | - | (section-grpo)= @@ -138,6 +140,7 @@ A dummy place holder of GRPO config for backward compatibility. | `evaluator` | [`EvaluatorConfig`](section-evaluator) | **Required** | - | | `stats_logger` | [`StatsLoggerConfig`](section-stats-logger) | **Required** | - | | `perf_tracer` | [`PerfTracerConfig`](section-perf-tracer) \| None | `None` | Performance tracer configuration. None means disabled. | +| `memory_profiler` | [`MemoryProfilerConfig`](section-memory-profiler) \| None | `None` | Memory snapshot profiler configuration. None means disabled. | | `recover` | [`RecoverConfig`](section-recover) | **Required** | - | | `sglang` | [`SGLangConfig`](section-sg-lang) | **Required** | - | | `vllm` | [`vLLMConfig`](section-v-llm) | **Required** | - | @@ -175,6 +178,7 @@ Configuration for Proximal Policy Optimization (PPO) reinforcement learning expe | `evaluator` | [`EvaluatorConfig`](section-evaluator) | **Required** | - | | `stats_logger` | [`StatsLoggerConfig`](section-stats-logger) | **Required** | - | | `perf_tracer` | [`PerfTracerConfig`](section-perf-tracer) \| None | `None` | Performance tracer configuration. None means disabled. | +| `memory_profiler` | [`MemoryProfilerConfig`](section-memory-profiler) \| None | `None` | Memory snapshot profiler configuration. None means disabled. | | `recover` | [`RecoverConfig`](section-recover) | **Required** | - | | `sglang` | [`SGLangConfig`](section-sg-lang) | **Required** | - | | `vllm` | [`vLLMConfig`](section-v-llm) | **Required** | - | @@ -194,29 +198,30 @@ Configuration for Proximal Policy Optimization (PPO) reinforcement learning expe Configuration for Reward Model (RW) training experiments. -| Parameter | Type | Default | Description | -| -------------------- | ----------------------------------------------------- | ------------ | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | -| `experiment_name` | string | **Required** | Name of the experiment (no '\_' or '/'). Required. | -| `trial_name` | string | **Required** | Name of the trial (no '-' or '/'). Required. | -| `cluster` | [`ClusterSpecConfig`](section-cluster) | **Required** | Cluster specification. Mainly used by slurm. | -| `allocation_mode` | string | `""` | DEPRECATED: Use per-engine 'backend' fields instead (e.g., actor.backend, rollout.backend). Legacy pattern-based GPU parallel strategy allocation mode. Only used by SPMD launchers (local/ray/slurm). Manual migration to per-engine 'backend' fields is required. | -| `seed` | integer | `1` | Random seed for reproducibility. | -| `enable_offload` | boolean | `False` | Whether to enable training offload using torch_memory_saver. This requires setting up the environment for TMS (e.g., via LD_PRELOAD). | -| `total_train_epochs` | integer | `1` | Total number of epochs to train the model. | -| `total_train_steps` | integer \| None | `None` | Terminate training after this number of steps. For benchmarking purposes only. None indicates normal training. | -| `total_train_n_seqs` | integer \| None | `None` | Terminate training after consuming this number of samples. For benchmarking purposes only. None indicates normal training. | -| `tokenizer_path` | string | `""` | Path to the tokenizer. | -| `train_dataset` | [`TrainDatasetConfig`](section-train-dataset) | **Required** | - | -| `valid_dataset` | [`ValidDatasetConfig`](section-valid-dataset) \| None | `None` | - | -| `saver` | [`SaverConfig`](section-saver) | **Required** | - | -| `evaluator` | [`EvaluatorConfig`](section-evaluator) | **Required** | - | -| `stats_logger` | [`StatsLoggerConfig`](section-stats-logger) | **Required** | - | -| `perf_tracer` | [`PerfTracerConfig`](section-perf-tracer) \| None | `None` | Performance tracer configuration. None means disabled. | -| `recover` | [`RecoverConfig`](section-recover) | **Required** | - | -| `sglang` | [`SGLangConfig`](section-sg-lang) | **Required** | - | -| `vllm` | [`vLLMConfig`](section-v-llm) | **Required** | - | -| `scheduler` | [`SchedulerConfig`](section-scheduler) | **Required** | - | -| `actor` | [`TrainEngineConfig`](section-train-engine) | **Required** | - | +| Parameter | Type | Default | Description | +| -------------------- | --------------------------------------------------------- | ------------ | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | +| `experiment_name` | string | **Required** | Name of the experiment (no '\_' or '/'). Required. | +| `trial_name` | string | **Required** | Name of the trial (no '-' or '/'). Required. | +| `cluster` | [`ClusterSpecConfig`](section-cluster) | **Required** | Cluster specification. Mainly used by slurm. | +| `allocation_mode` | string | `""` | DEPRECATED: Use per-engine 'backend' fields instead (e.g., actor.backend, rollout.backend). Legacy pattern-based GPU parallel strategy allocation mode. Only used by SPMD launchers (local/ray/slurm). Manual migration to per-engine 'backend' fields is required. | +| `seed` | integer | `1` | Random seed for reproducibility. | +| `enable_offload` | boolean | `False` | Whether to enable training offload using torch_memory_saver. This requires setting up the environment for TMS (e.g., via LD_PRELOAD). | +| `total_train_epochs` | integer | `1` | Total number of epochs to train the model. | +| `total_train_steps` | integer \| None | `None` | Terminate training after this number of steps. For benchmarking purposes only. None indicates normal training. | +| `total_train_n_seqs` | integer \| None | `None` | Terminate training after consuming this number of samples. For benchmarking purposes only. None indicates normal training. | +| `tokenizer_path` | string | `""` | Path to the tokenizer. | +| `train_dataset` | [`TrainDatasetConfig`](section-train-dataset) | **Required** | - | +| `valid_dataset` | [`ValidDatasetConfig`](section-valid-dataset) \| None | `None` | - | +| `saver` | [`SaverConfig`](section-saver) | **Required** | - | +| `evaluator` | [`EvaluatorConfig`](section-evaluator) | **Required** | - | +| `stats_logger` | [`StatsLoggerConfig`](section-stats-logger) | **Required** | - | +| `perf_tracer` | [`PerfTracerConfig`](section-perf-tracer) \| None | `None` | Performance tracer configuration. None means disabled. | +| `memory_profiler` | [`MemoryProfilerConfig`](section-memory-profiler) \| None | `None` | Memory snapshot profiler configuration. None means disabled. | +| `recover` | [`RecoverConfig`](section-recover) | **Required** | - | +| `sglang` | [`SGLangConfig`](section-sg-lang) | **Required** | - | +| `vllm` | [`vLLMConfig`](section-v-llm) | **Required** | - | +| `scheduler` | [`SchedulerConfig`](section-scheduler) | **Required** | - | +| `actor` | [`TrainEngineConfig`](section-train-engine) | **Required** | - | (section-sft)= @@ -224,29 +229,30 @@ Configuration for Reward Model (RW) training experiments. Configuration for Supervised Fine-Tuning (SFT) experiments. -| Parameter | Type | Default | Description | -| -------------------- | ----------------------------------------------------- | ------------ | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | -| `experiment_name` | string | **Required** | Name of the experiment (no '\_' or '/'). Required. | -| `trial_name` | string | **Required** | Name of the trial (no '-' or '/'). Required. | -| `cluster` | [`ClusterSpecConfig`](section-cluster) | **Required** | Cluster specification. Mainly used by slurm. | -| `allocation_mode` | string | `""` | DEPRECATED: Use per-engine 'backend' fields instead (e.g., actor.backend, rollout.backend). Legacy pattern-based GPU parallel strategy allocation mode. Only used by SPMD launchers (local/ray/slurm). Manual migration to per-engine 'backend' fields is required. | -| `seed` | integer | `1` | Random seed for reproducibility. | -| `enable_offload` | boolean | `False` | Whether to enable training offload using torch_memory_saver. This requires setting up the environment for TMS (e.g., via LD_PRELOAD). | -| `total_train_epochs` | integer | `1` | Total number of epochs to train the model. | -| `total_train_steps` | integer \| None | `None` | Terminate training after this number of steps. For benchmarking purposes only. None indicates normal training. | -| `total_train_n_seqs` | integer \| None | `None` | Terminate training after consuming this number of samples. For benchmarking purposes only. None indicates normal training. | -| `tokenizer_path` | string | `""` | Path to the tokenizer. | -| `train_dataset` | [`TrainDatasetConfig`](section-train-dataset) | **Required** | - | -| `valid_dataset` | [`ValidDatasetConfig`](section-valid-dataset) \| None | `None` | - | -| `saver` | [`SaverConfig`](section-saver) | **Required** | - | -| `evaluator` | [`EvaluatorConfig`](section-evaluator) | **Required** | - | -| `stats_logger` | [`StatsLoggerConfig`](section-stats-logger) | **Required** | - | -| `perf_tracer` | [`PerfTracerConfig`](section-perf-tracer) \| None | `None` | Performance tracer configuration. None means disabled. | -| `recover` | [`RecoverConfig`](section-recover) | **Required** | - | -| `sglang` | [`SGLangConfig`](section-sg-lang) | **Required** | - | -| `vllm` | [`vLLMConfig`](section-v-llm) | **Required** | - | -| `scheduler` | [`SchedulerConfig`](section-scheduler) | **Required** | - | -| `actor` | [`TrainEngineConfig`](section-train-engine) | **Required** | - | +| Parameter | Type | Default | Description | +| -------------------- | --------------------------------------------------------- | ------------ | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | +| `experiment_name` | string | **Required** | Name of the experiment (no '\_' or '/'). Required. | +| `trial_name` | string | **Required** | Name of the trial (no '-' or '/'). Required. | +| `cluster` | [`ClusterSpecConfig`](section-cluster) | **Required** | Cluster specification. Mainly used by slurm. | +| `allocation_mode` | string | `""` | DEPRECATED: Use per-engine 'backend' fields instead (e.g., actor.backend, rollout.backend). Legacy pattern-based GPU parallel strategy allocation mode. Only used by SPMD launchers (local/ray/slurm). Manual migration to per-engine 'backend' fields is required. | +| `seed` | integer | `1` | Random seed for reproducibility. | +| `enable_offload` | boolean | `False` | Whether to enable training offload using torch_memory_saver. This requires setting up the environment for TMS (e.g., via LD_PRELOAD). | +| `total_train_epochs` | integer | `1` | Total number of epochs to train the model. | +| `total_train_steps` | integer \| None | `None` | Terminate training after this number of steps. For benchmarking purposes only. None indicates normal training. | +| `total_train_n_seqs` | integer \| None | `None` | Terminate training after consuming this number of samples. For benchmarking purposes only. None indicates normal training. | +| `tokenizer_path` | string | `""` | Path to the tokenizer. | +| `train_dataset` | [`TrainDatasetConfig`](section-train-dataset) | **Required** | - | +| `valid_dataset` | [`ValidDatasetConfig`](section-valid-dataset) \| None | `None` | - | +| `saver` | [`SaverConfig`](section-saver) | **Required** | - | +| `evaluator` | [`EvaluatorConfig`](section-evaluator) | **Required** | - | +| `stats_logger` | [`StatsLoggerConfig`](section-stats-logger) | **Required** | - | +| `perf_tracer` | [`PerfTracerConfig`](section-perf-tracer) \| None | `None` | Performance tracer configuration. None means disabled. | +| `memory_profiler` | [`MemoryProfilerConfig`](section-memory-profiler) \| None | `None` | Memory snapshot profiler configuration. None means disabled. | +| `recover` | [`RecoverConfig`](section-recover) | **Required** | - | +| `sglang` | [`SGLangConfig`](section-sg-lang) | **Required** | - | +| `vllm` | [`vLLMConfig`](section-v-llm) | **Required** | - | +| `scheduler` | [`SchedulerConfig`](section-scheduler) | **Required** | - | +| `actor` | [`TrainEngineConfig`](section-train-engine) | **Required** | - | (section-fsdp-engine)= @@ -968,6 +974,20 @@ Refer to Megatron-LM documentation for implementation details. | `fp8_config` | [`FP8EngineConfig`](section-fp8-engine) \| None | `None` | - | | `bridge_type` | string | `"mbridge"` | Bridge backend for MegatronEngine. Choices: 'mbridge' or 'megatron-bridge'. **Choices:** `mbridge`, `megatron-bridge` | +(section-memory-profiler)= + +## MemoryProfiler Configuration + +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. + +| Parameter | Type | Default | Description | +| --------------- | --------------- | ------------ | ------------------------------------------------- | +| `profile_steps` | list of integer | **Required** | List of global steps to capture memory snapshots. | +| `max_entries` | integer | `100000` | Max entries for memory history ring buffer. | + (section-open-ai-proxy)= ## OpenAIProxy Configuration