Configuration

@dataclass
class SkyRLTrainConfig(BaseConfig):
    data: DataConfig = field(default_factory=DataConfig)
    trainer: TrainerConfig = field(default_factory=TrainerConfig)
    generator: GeneratorConfig = field(default_factory=GeneratorConfig)
    environment: EnvironmentConfig = field(default_factory=EnvironmentConfig)

    def __post_init__(self):

        # generator.max_input_length defaults to trainer.max_prompt_length
        if self.generator.max_input_length is None:
            self.generator.max_input_length = self.trainer.max_prompt_length

        # Copy temperature from generator sampling params to algorithm config
        # so workers can access it without needing the generator config
        if self.trainer.algorithm.temperature is None:
            self.trainer.algorithm.temperature = self.generator.sampling_params.temperature

        if self.data.dataloader.num_workers is None:
            self.data.dataloader.num_workers = 8
        if self.data.dataloader.persistent_workers and self.data.dataloader.num_workers == 0:
            raise ValueError(
                "data.dataloader.persistent_workers requires num_workers > 0, but it was set explicitly to 0."
            )

        # TODO(devpatel): Bandaid solution, replace this once we have a better
        # solution for LoRA performance degradation on the vLLM side
        from skyrl.backends.skyrl_train.inference_servers.utils import (
            _uses_lora_weight_sync,
        )

        ie_cfg = self.generator.inference_engine
        if _uses_lora_weight_sync(self) and ie_cfg.enforce_eager and ie_cfg.backend == "vllm":
            import warnings

            warnings.warn(
                "LoRA is enabled but inference_engine.enforce_eager=true. "
                "This combination causes significant performance degradation (2-3x slower generation). "
                "Automatically setting enforce_eager=false for better performance. "
            )
            ie_cfg.enforce_eager = False

    @classmethod
    def from_cli_overrides(cls, args: Union[List[str], dict]) -> "SkyRLTrainConfig":
        """Construct a SkyRLTrainConfig from CLI arguments or a dict of overrides.

        Parses CLI arguments and builds a typed config. Dataclass field defaults
        are used for any values not specified on the command line.

        Args:
            args: Either a list of CLI arguments in 'key.path=value' format, or a dict
                  mapping dot-notation keys to values.
                  Example list: ['trainer.policy.model.path=Qwen/Qwen2.5-1.5B-Instruct', 'trainer.seed=123']
                  Example dict: {'trainer.policy.model.path': 'Qwen/Qwen2.5-1.5B-Instruct', 'trainer.seed': 123}

        Returns:
            A fully constructed SkyRLTrainConfig with CLI overrides applied.

        Raises:
            ValueError: If an argument uses the unsupported '+' prefix.
        """
        if isinstance(args, dict):
            # OmegaConf's CLI parser only treats "null" as None; Python's
            # None stringifies to "None" which is parsed as the literal
            # string. Map None -> "null" so JSON-style overrides survive
            # the round-trip through OmegaConf.from_cli below.
            args = [f"{k}=null" if v is None else f"{k}={v}" for k, v in args.items()]

        # Check for unsupported '+' prefix
        for arg in args:
            if arg.startswith("+"):
                raise ValueError(
                    f"The '+' prefix for adding new config fields is not supported: '{arg}'. "
                    "To add custom config fields, subclass the relevant config dataclass."
                )
        overrides = OmegaConf.from_cli(args)
        unsupported_rope_paths = (
            "trainer.rope_scaling",
            "trainer.rope_theta",
            "trainer.rope_parameters",
            "generator.rope_scaling",
            "generator.rope_theta",
            "generator.rope_parameters",
            "generator.inference_engine.rope_scaling",
            "generator.inference_engine.rope_theta",
            "generator.inference_engine.rope_parameters",
            "generator.inference_engine.engine_init_kwargs.rope_scaling",
            "generator.inference_engine.engine_init_kwargs.rope_theta",
            "generator.inference_engine.engine_init_kwargs.rope_parameters",
            "generator.inference_engine.engine_init_kwargs.hf_overrides.rope_scaling",
            "generator.inference_engine.engine_init_kwargs.hf_overrides.rope_theta",
        )
        if any(_has_nested_key(overrides, path) for path in unsupported_rope_paths):
            raise ValueError(
                "`rope_scaling`, `rope_theta`, and `rope_parameters` are no longer supported as native "
                "config overrides, use `generator.inference_engine.engine_init_kwargs.hf_overrides.rope_parameters` "
                "and `trainer.policy.model_config_kwargs.rope_parameters` or "
                "`trainer.policy.megatron_config.transformer_config_kwargs.rope_parameters` instead"
            )
        inference_rope_parameters = _get_nested_value(
            overrides, "generator.inference_engine.engine_init_kwargs.hf_overrides.rope_parameters"
        )
        if inference_rope_parameters is not _MISSING:
            trainer_strategy = _get_nested_value(overrides, "trainer.strategy")
            trainer_strategy = "fsdp" if trainer_strategy is _MISSING else trainer_strategy
            trainer_rope_parameters_path = (
                "trainer.policy.megatron_config.transformer_config_kwargs.rope_parameters"
                if trainer_strategy == "megatron"
                else "trainer.policy.model_config_kwargs.rope_parameters"
            )
            trainer_rope_parameters = _get_nested_value(overrides, trainer_rope_parameters_path)
            if inference_rope_parameters != trainer_rope_parameters:
                raise ValueError(
                    "`generator.inference_engine.engine_init_kwargs.hf_overrides.rope_parameters` must match "
                    f"the trainer-side override at `{trainer_rope_parameters_path}`"
                )
        async_engine_path = "generator.inference_engine.async_engine"
        async_engine = _get_nested_value(overrides, async_engine_path)
        if async_engine is not _MISSING:
            if async_engine is True or (isinstance(async_engine, str) and async_engine.lower() == "true"):
                _delete_nested_key(overrides, async_engine_path)
            elif async_engine is False or (isinstance(async_engine, str) and async_engine.lower() == "false"):
                raise ValueError(
                    "`async_engine=False` is no longer supported; SkyRL always uses the async "
                    "HTTP/vLLM inference path. Remove the override."
                )
            else:
                raise ValueError("`async_engine` is no longer supported as a config field. Remove the override.")
        removed_inference_engine_overrides = {
            "generator.inference_engine.enable_http_endpoint": (
                "`enable_http_endpoint` is no longer supported; SkyRL always uses the HTTP/vLLM inference path. "
                "Remove the override."
            ),
            "generator.inference_engine.override_existing_update_group": (
                "`override_existing_update_group` is no longer supported; update-group handling is managed "
                "automatically by the vLLM-native inference path. Remove the override."
            ),
        }
        for path, message in removed_inference_engine_overrides.items():
            if _has_nested_key(overrides, path):
                raise ValueError(message)
        if (
            "generator" in overrides
            and "inference_engine" in overrides.generator
            and "remote_urls" in overrides.generator.inference_engine
        ):
            raise ValueError(
                "`remote_urls` is no longer supported, external inference servers can be used with "
                "`external_proxy_url` and `external_server_urls` instead"
            )
        # Accept the deprecated ``trainer.use_sample_packing`` key as an alias
        # for ``trainer.remove_microbatch_padding``. Remap it before
        # construction so the strict key validation does not reject the old
        # name.
        if "trainer" in overrides and "use_sample_packing" in overrides.trainer:
            if "remove_microbatch_padding" in overrides.trainer:
                raise ValueError(
                    "Specify only one of trainer.use_sample_packing (deprecated) and "
                    "trainer.remove_microbatch_padding, not both."
                )
            import warnings

            warnings.warn(
                "trainer.use_sample_packing has been renamed to "
                "trainer.remove_microbatch_padding; use "
                "trainer.remove_microbatch_padding instead.",
                DeprecationWarning,
                stacklevel=2,
            )
            overrides.trainer["remove_microbatch_padding"] = overrides.trainer["use_sample_packing"]
            del overrides.trainer["use_sample_packing"]
        return cls.from_dict_config(overrides)

    @classmethod
    def from_cli_overrides(cls, args: Union[List[str], dict]) -> "SkyRLTrainConfig":
        """Construct a SkyRLTrainConfig from CLI arguments or a dict of overrides.

        Parses CLI arguments and builds a typed config. Dataclass field defaults
        are used for any values not specified on the command line.

        Args:
            args: Either a list of CLI arguments in 'key.path=value' format, or a dict
                  mapping dot-notation keys to values.
                  Example list: ['trainer.policy.model.path=Qwen/Qwen2.5-1.5B-Instruct', 'trainer.seed=123']
                  Example dict: {'trainer.policy.model.path': 'Qwen/Qwen2.5-1.5B-Instruct', 'trainer.seed': 123}

        Returns:
            A fully constructed SkyRLTrainConfig with CLI overrides applied.

        Raises:
            ValueError: If an argument uses the unsupported '+' prefix.
        """
        if isinstance(args, dict):
            # OmegaConf's CLI parser only treats "null" as None; Python's
            # None stringifies to "None" which is parsed as the literal
            # string. Map None -> "null" so JSON-style overrides survive
            # the round-trip through OmegaConf.from_cli below.
            args = [f"{k}=null" if v is None else f"{k}={v}" for k, v in args.items()]

        # Check for unsupported '+' prefix
        for arg in args:
            if arg.startswith("+"):
                raise ValueError(
                    f"The '+' prefix for adding new config fields is not supported: '{arg}'. "
                    "To add custom config fields, subclass the relevant config dataclass."
                )
        overrides = OmegaConf.from_cli(args)
        unsupported_rope_paths = (
            "trainer.rope_scaling",
            "trainer.rope_theta",
            "trainer.rope_parameters",
            "generator.rope_scaling",
            "generator.rope_theta",
            "generator.rope_parameters",
            "generator.inference_engine.rope_scaling",
            "generator.inference_engine.rope_theta",
            "generator.inference_engine.rope_parameters",
            "generator.inference_engine.engine_init_kwargs.rope_scaling",
            "generator.inference_engine.engine_init_kwargs.rope_theta",
            "generator.inference_engine.engine_init_kwargs.rope_parameters",
            "generator.inference_engine.engine_init_kwargs.hf_overrides.rope_scaling",
            "generator.inference_engine.engine_init_kwargs.hf_overrides.rope_theta",
        )
        if any(_has_nested_key(overrides, path) for path in unsupported_rope_paths):
            raise ValueError(
                "`rope_scaling`, `rope_theta`, and `rope_parameters` are no longer supported as native "
                "config overrides, use `generator.inference_engine.engine_init_kwargs.hf_overrides.rope_parameters` "
                "and `trainer.policy.model_config_kwargs.rope_parameters` or "
                "`trainer.policy.megatron_config.transformer_config_kwargs.rope_parameters` instead"
            )
        inference_rope_parameters = _get_nested_value(
            overrides, "generator.inference_engine.engine_init_kwargs.hf_overrides.rope_parameters"
        )
        if inference_rope_parameters is not _MISSING:
            trainer_strategy = _get_nested_value(overrides, "trainer.strategy")
            trainer_strategy = "fsdp" if trainer_strategy is _MISSING else trainer_strategy
            trainer_rope_parameters_path = (
                "trainer.policy.megatron_config.transformer_config_kwargs.rope_parameters"
                if trainer_strategy == "megatron"
                else "trainer.policy.model_config_kwargs.rope_parameters"
            )
            trainer_rope_parameters = _get_nested_value(overrides, trainer_rope_parameters_path)
            if inference_rope_parameters != trainer_rope_parameters:
                raise ValueError(
                    "`generator.inference_engine.engine_init_kwargs.hf_overrides.rope_parameters` must match "
                    f"the trainer-side override at `{trainer_rope_parameters_path}`"
                )
        async_engine_path = "generator.inference_engine.async_engine"
        async_engine = _get_nested_value(overrides, async_engine_path)
        if async_engine is not _MISSING:
            if async_engine is True or (isinstance(async_engine, str) and async_engine.lower() == "true"):
                _delete_nested_key(overrides, async_engine_path)
            elif async_engine is False or (isinstance(async_engine, str) and async_engine.lower() == "false"):
                raise ValueError(
                    "`async_engine=False` is no longer supported; SkyRL always uses the async "
                    "HTTP/vLLM inference path. Remove the override."
                )
            else:
                raise ValueError("`async_engine` is no longer supported as a config field. Remove the override.")
        removed_inference_engine_overrides = {
            "generator.inference_engine.enable_http_endpoint": (
                "`enable_http_endpoint` is no longer supported; SkyRL always uses the HTTP/vLLM inference path. "
                "Remove the override."
            ),
            "generator.inference_engine.override_existing_update_group": (
                "`override_existing_update_group` is no longer supported; update-group handling is managed "
                "automatically by the vLLM-native inference path. Remove the override."
            ),
        }
        for path, message in removed_inference_engine_overrides.items():
            if _has_nested_key(overrides, path):
                raise ValueError(message)
        if (
            "generator" in overrides
            and "inference_engine" in overrides.generator
            and "remote_urls" in overrides.generator.inference_engine
        ):
            raise ValueError(
                "`remote_urls` is no longer supported, external inference servers can be used with "
                "`external_proxy_url` and `external_server_urls` instead"
            )
        # Accept the deprecated ``trainer.use_sample_packing`` key as an alias
        # for ``trainer.remove_microbatch_padding``. Remap it before
        # construction so the strict key validation does not reject the old
        # name.
        if "trainer" in overrides and "use_sample_packing" in overrides.trainer:
            if "remove_microbatch_padding" in overrides.trainer:
                raise ValueError(
                    "Specify only one of trainer.use_sample_packing (deprecated) and "
                    "trainer.remove_microbatch_padding, not both."
                )
            import warnings

            warnings.warn(
                "trainer.use_sample_packing has been renamed to "
                "trainer.remove_microbatch_padding; use "
                "trainer.remove_microbatch_padding instead.",
                DeprecationWarning,
                stacklevel=2,
            )
            overrides.trainer["remove_microbatch_padding"] = overrides.trainer["use_sample_packing"]
            del overrides.trainer["use_sample_packing"]
        return cls.from_dict_config(overrides)

@dataclass
class DataConfig(BaseConfig):
    train_data: List[str] = field(default_factory=lambda: [os.path.expanduser("~/data/gsm8k/train.parquet")])
    val_data: List[str] = field(default_factory=lambda: [os.path.expanduser("~/data/gsm8k/validation.parquet")])
    dataloader: DataLoaderConfig = field(default_factory=DataLoaderConfig)

@dataclass
class ModelConfig(BaseConfig):
    path: Optional[str] = None
    lora: SkyRLLoraConfig = field(default_factory=SkyRLLoraConfig)
    fake_int4_qat: FakeInt4QatConfig = field(default_factory=FakeInt4QatConfig)

    def __post_init__(self) -> None:
        if self.fake_int4_qat.enabled:
            assert self.lora.rank > 0, (
                "`trainer.policy.model.fake_int4_qat.enabled=True` currently requires LoRA "
                "(`trainer.policy.model.lora.rank > 0`) because full-weight sync exports "
                "dense expert weights."
            )

@dataclass
class SkyRLLoraConfig(BaseConfig):
    rank: int = 0
    alpha: int = 16
    dropout: float = 0.0
    lora_sync_path: str = "/tmp/skyrl_lora_sync"
    target_modules: str = "all-linear"
    exclude_modules: Optional[str] = None
    init_method: str = "kaiming"
    """For FSDP, corresponds to ``init_lora_weights`` in PEFT.
    For Megatron, used for ``lora_A_init_method``; supports "xavier", "normal", "kaiming", "zero"."""

    max_loras: int = 1
    """Maximum number of LoRA adapters that can be active concurrently in a
    single GPU batch. Maps to vLLM's ``max_loras``. Increase past 1 to enable
    multi-tenant LoRA serving via ``RemoteInferenceClient.load_lora_adapter``."""

    max_cpu_loras: Optional[int] = None
    """Total LoRA adapter capacity in vLLM's CPU LRU cache. Maps to vLLM's
    ``max_cpu_loras``; when None, vLLM defaults it to ``max_loras``. Must be
    >= ``max_loras`` if explicitly set."""

@dataclass
class TrainerConfig(BaseConfig):
    placement: PlacementConfig = field(default_factory=PlacementConfig)
    use_expandable_segments: bool = True
    """Enable PyTorch's CUDA ``expandable_segments`` allocator on the training
    workers to reduce GPU memory fragmentation across the offload/backload and
    forward/backward cycles. See ``InferenceEngineConfig`` for the
    equivalent inference-engine knob."""
    sequence_parallel_backend: str = "ulysses"
    strategy: str = "fsdp"
    policy: PolicyConfig = field(default_factory=PolicyConfig)
    ref: RefConfig = field(default_factory=RefConfig)
    critic: CriticConfig = field(default_factory=CriticConfig)
    algorithm: AlgorithmConfig = field(default_factory=AlgorithmConfig)
    mtp: MTPConfig = field(default_factory=MTPConfig)
    fully_async: FullyAsyncConfig = field(default_factory=FullyAsyncConfig)
    gradient_checkpointing: bool = True
    gradient_checkpointing_use_reentrant: bool = False
    seed: int = 42
    resume_mode: Optional[str] = "latest"
    """``None``/``"none"``, ``"latest"``, or ``"from_path"``."""
    resume_path: Optional[str] = None
    log_path: str = "/tmp/skyrl-logs"
    """Path for infrastructure log files. For multi-node, use a shared filesystem path to consolidate logs."""
    ckpt_path: str = field(default_factory=lambda: os.path.expanduser("~/ckpts/"))
    max_ckpts_to_keep: int = -1
    """``-1`` to keep all checkpoints, ``N`` to keep only the last N."""
    ckpt_interval: int = 10
    hf_save_interval: int = -1
    """Save HuggingFace-format model every N steps. ``-1`` to disable."""
    export_path: str = field(default_factory=lambda: os.path.expanduser("~/exports/"))
    """Path for exported artifacts (HF models, debug dumps, etc.)."""
    bf16: bool = True
    epochs: int = 1
    max_training_steps: Optional[int] = None
    """If set, stop training after this many steps regardless of epochs or dataset size.
    Useful for CI smoke tests and quick validation runs."""
    update_epochs_per_batch: int = 1
    """Number of gradient update passes over each training batch."""
    train_batch_size: int = 1024
    """See ``utils/utils.py::validate_batch_sizes`` for train, mini, and micro batch size constraints."""
    policy_mini_batch_size: int = 256
    critic_mini_batch_size: int = 256
    micro_train_batch_size_per_gpu: int = 1
    micro_forward_batch_size_per_gpu: int = 1
    max_tokens_per_microbatch: int = -1
    """Maximum number of tokens per microbatch for both forward and training steps. When > 0, microbatches 
    are formed by bin-packing samples based on their token counts (from attention_mask) instead of using a 
    fixed sample count, and micro_train_batch_size_per_gpu / micro_forward_batch_size_per_gpu are ignored.
    -1 means disabled (use sample-based micro_train_batch_size_per_gpu / micro_forward_batch_size_per_gpu).
    Applies to both forward and training micro-batching.

    NOTE: this is a *soft* cap. Sequences are never split across microbatches, so a single sequence
    longer than ``max_tokens_per_microbatch`` is placed alone in its own microbatch that exceeds the
    cap (no error, no truncation). The true peak microbatch size is therefore
    ``max(max_tokens_per_microbatch, longest_sequence_in_batch)``."""
    recompute_old_logprobs_per_minibatch: bool = True
    """When True, recomputes policy/ref model logprobs (and critic values) per mini-batch using
    the same mini-batch + DP partition as the training step. When False, a single full-batch forward is run.
    This makes the microbatch packing — and therefore the resulting logprobs/values — identical to
    what forward_backward recomputes, so the PPO ratio (and critic value clipping) is exact at the
    first inner step."""
    update_ref_every_epoch: bool = False
    remove_microbatch_padding: bool = True
    """Pack samples into the THD layout and strip intra-microbatch padding (requires flash attention)."""
    eval_batch_size: int = 1024
    eval_before_train: bool = True
    eval_interval: int = 5
    """``-1`` to disable evaluation."""
    max_prompt_length: int = 512
    flash_attn: bool = True
    disable_fast_tokenizer: bool = False
    project_name: str = "skyrl"
    run_name: str = "test_run"
    logger: str = "wandb"
    enable_ray_gpu_monitor: bool = True
    """Enable background Ray GPU/RAM metrics collection and logging to wandb."""
    tags: Optional[List[str]] = None
    """Optional list of tags to apply to the W&B run. Has no effect on other backends."""
    dump_data_batch: bool = False
    dump_eval_results: bool = True
    print_example_interval: int = 1
    """Pretty-print an example prompt/response/reward to stdout every N
    training steps; ``0``/``-1`` disables. Renamed from ``log_example_interval``."""
    num_logger_eval_samples: int = -1
    """Number of evaluation trajectory (prompt, response, score) tuples to upload to a wandb
    table on each eval. ``-1`` (default) or ``0`` disables. When positive,
    up to this many samples are taken from the start of each eval pass and
    logged via :class:`TrajectoryLogger`. Column count is fixed
    by the first call, so keep the eval set size and this value stable."""
    num_logger_train_samples: int = -1
    """Number of training trajectory (prompt, response, score) tuples to upload to a wandb
    table on each training step. ``-1`` (default) or ``0`` disables. When positive,
    up to this many samples are taken from the start of each training step and
    logged via :class:`TrajectoryLogger`. Column count is fixed
    by the first call, so keep the training set size and this value stable."""
    log_example_interval: int = -1
    """Log an example prompt every N training steps, ``0``/``-1`` to disable"""
    logprobs_chunk_size: Optional[int] = 1024
    """Chunk size along the sequence dimension when computing log-probs from logits.
    This lowers peak GPU memory at the cost of ~2x wall-clock time.
    ``None`` disables chunking (Megatron backend only; FSDP requires a positive int).
    See https://github.com/NovaSky-AI/SkyRL/pull/1610 for more details."""
    vocab_entropy_chunk_size: Optional[int] = 0
    """Chunk size along the sequence dimension when computing Megatron vocab entropy.
    ``0`` auto-sizes from the local vocab shard size and ``vocab_entropy_chunk_memory_mb``.
    ``None`` disables chunking."""
    vocab_entropy_chunk_memory_mb: int = 512
    """Approximate per-chunk temporary memory budget for auto-sized Megatron vocab entropy chunks."""
    fused_lm_head_logprob: bool = False
    """Megatron only. Fuse the LM-head projection into log-prob / entropy
    computation so the full ``[B, S, vocab//TP]`` logits tensor is never
    materialized. Uses ``logprobs_chunk_size`` to bound peak memory."""
    fused_lm_head_logprob_backend: str = "torch"
    """Fused LM-head backend: ``"torch"`` (default) or ``"triton"``.
    The Triton backend requires CUDA + triton and falls back to ``"torch"``
    when unavailable. Ignored unless ``fused_lm_head_logprob`` is true."""

    def __post_init__(self):
        # ref model defaults to the policy model
        if self.ref.model.path is None:
            self.ref.model.path = self.policy.model.path

        if self.log_example_interval > 0:
            print(
                f"log_example_interval has been renamed, use print_example_interval instead. Setting print_example_interval to {self.log_example_interval}"
            )
            self.print_example_interval = self.log_example_interval

        if self.policy.model.fake_int4_qat.enabled:
            assert (
                self.strategy == "megatron"
            ), "`trainer.policy.model.fake_int4_qat.enabled=True` is only supported with `trainer.strategy=megatron`."
            assert not self.policy.megatron_config.lora_config.merge_lora, (
                "`trainer.policy.model.fake_int4_qat.enabled=True` currently requires "
                "`trainer.policy.megatron_config.lora_config.merge_lora=False` so weight "
                "sync preserves the inference engine's INT4 base weights."
            )

        if self.logprobs_chunk_size is not None and (
            not isinstance(self.logprobs_chunk_size, int) or self.logprobs_chunk_size <= 0
        ):
            raise ValueError(
                f"logprobs_chunk_size must be a positive integer or None, got {self.logprobs_chunk_size!r}."
            )
        if self.logprobs_chunk_size is None and self.strategy != "megatron":
            raise ValueError(
                "logprobs_chunk_size=None (no chunking) is only supported with the Megatron backend. "
                f"Set a positive integer for strategy={self.strategy!r}."
            )
        if self.fused_lm_head_logprob and self.strategy != "megatron":
            raise ValueError(
                "fused_lm_head_logprob=True is only supported with the Megatron backend, "
                f"got strategy={self.strategy!r}."
            )
        if self.fused_lm_head_logprob_backend not in ("torch", "triton"):
            raise ValueError(
                "fused_lm_head_logprob_backend must be 'torch' or 'triton', "
                f"got {self.fused_lm_head_logprob_backend!r}."
            )
        if self.vocab_entropy_chunk_size is not None and (
            isinstance(self.vocab_entropy_chunk_size, bool)
            or not isinstance(self.vocab_entropy_chunk_size, int)
            or self.vocab_entropy_chunk_size < 0
        ):
            raise ValueError(
                "vocab_entropy_chunk_size must be a non-negative integer or None, "
                f"got {self.vocab_entropy_chunk_size!r}."
            )
        if (
            isinstance(self.vocab_entropy_chunk_memory_mb, bool)
            or not isinstance(self.vocab_entropy_chunk_memory_mb, int)
            or self.vocab_entropy_chunk_memory_mb <= 0
        ):
            raise ValueError(
                "vocab_entropy_chunk_memory_mb must be a positive integer, "
                f"got {self.vocab_entropy_chunk_memory_mb!r}."
            )

@dataclass
class OptimizerConfig(BaseConfig):
    lr: float = 1e-6
    adam_betas: List[float] = field(default_factory=lambda: [0.9, 0.999])
    weight_decay: float = 1e-2
    max_grad_norm: float = 1.0
    offload_after_step: bool = True
    """Offload optimizer state to CPU after each full training step. Only applicable when ``colocate_all=True``."""
    num_warmup_steps: int = 0
    """Number of mini-batch steps to warmup the optimizer."""
    scheduler: str = "constant_with_warmup"

@dataclass
class MixedPrecisionConfig(BaseConfig):
    param_dtype: str = "bf16"
    reduce_dtype: str = "fp32"
    buffer_dtype: str = "fp32"

@dataclass
class FSDPConfig(BaseConfig):
    cpu_offload: bool = False
    """Offload params and optimizer state to CPU during the forward pass."""
    reshard_after_forward: Union[bool, int] = True
    """FSDP2 only. Accepts True, False, or an int between 1 and ``fsdp_size``."""
    fsdp_size: int = -1
    mixed_precision: Optional[MixedPrecisionConfig] = None
    # specify wrap policy as a dict with `transformer_layer_cls_to_wrap` key for custom module based wrapping
    wrap_policy: dict = field(default_factory=dict)

@dataclass
class MegatronConfig(BaseConfig):
    tensor_model_parallel_size: int = 1
    pipeline_model_parallel_size: int = 1
    context_parallel_size: int = 1
    expert_model_parallel_size: int = 1
    expert_tensor_parallel_size: Optional[int] = None
    # MoE runtime configuration flags
    moe_token_dispatcher_type: str = "alltoall"
    moe_router_load_balancing_type: str = "none"
    """Set to "aux_loss", "seq_aux_loss", or "global_aux_loss" to enable aux loss-based load balancing and logging."""
    moe_aux_loss_coeff: float = 0.0
    """Scaling coefficient for the moe load balancing loss if moe_router_load_balancing_type is not 'none'. Will disable aux loss in megatron-core if set to 0."""
    moe_grouped_gemm: bool = True
    moe_router_score_function: Optional[str] = None
    moe_router_enable_expert_bias: Optional[bool] = None
    moe_enable_routing_replay: bool = False
    moe_per_layer_logging: bool = False
    """Enable per-layer logging of MoE metrics (i.e. per layer aux losses)."""
    moe_router_dtype: str = "fp32"
    """Pass through to Megatron-Bridge - can be set to 'fp64' for additional numerical stability."""
    ddp_config: MegatronDDPConfig = field(default_factory=MegatronDDPConfig)
    hf_export_config: MegatronHFExportConfig = field(default_factory=MegatronHFExportConfig)
    lora_config: MegatronLoraConfig = field(default_factory=MegatronLoraConfig)
    optimizer_config_kwargs: Dict[str, Any] = field(
        default_factory=lambda: copy.deepcopy(DEFAULT_MEGATRON_OPTIMIZER_KWARGS)
    )
    transformer_config_kwargs: Dict[str, Any] = field(
        default_factory=lambda: copy.deepcopy(DEFAULT_TRANSFORMER_CONFIG_KWARGS)
    )
    empty_cuda_cache: Optional[bool] = True
    model_config_kwargs: dict = field(default_factory=dict)
    dist_ckpt_optim_fully_reshardable: bool = False
    freeze_moe_router: bool = False
    """If True, freeze MoE router parameters so they are not updated during training. No-op on
    non-MoE models."""
    mtp_num_layers: Optional[int] = None
    """Number of Multi-Token Prediction (MTP) heads to build. ``None`` honors the model's HF config
    (``num_nextn_predict_layers``); an int overrides it (``0`` force-disables MTP). Active heads are
    trained with the decoupled draft loss and synced to vLLM for speculative decoding."""
    mtp_loss_weight: float = 0.1
    """Weight ``w`` of the draft loss in ``policy_loss + w * draft_loss``. The draft loss is fully
    decoupled: trunk, re-embedding, output weight and teacher are all detached, so its gradient
    reaches only the ``.mtp.`` head params. Only used when MTP heads are active."""
    mtp_loss_chunk_size: Optional[int] = 1024
    """Sequence-chunk size for the draft loss, with gradient checkpointing, to bound peak memory at
    large vocab (e.g. Qwen3.5's 248K, where the full-sequence softmax OOMs). Numerically identical to
    no chunking. ``None`` disables it; ignored when ``mtp_loss_topk`` is set."""
    mtp_loss_topk: Optional[int] = None
    """If set, use a top-k approximation of the soft-CE draft loss: distill only the teacher's top-k
    tokens (renormalized), ``O(seq*k)`` memory instead of ``O(seq*vocab)`` -- fits at large vocab
    without fragmentation. Reconciled across the TP group, so it scales to any parallel size.
    ``None`` uses the exact full-vocab loss. Typical: 64-128."""
    async_dist_ckpt_save: bool = False
    """Write the torch_dist checkpoint from a background process so training resumes
    immediately; the pending write is finalized at the next checkpoint and at shutdown.
    The on-disk format is identical to a synchronous save. Only the sharded
    model/optimizer state is async -- the rank-0 HF config/tokenizer write stays inline.
    Falls back to synchronous for cloud paths."""
    async_dist_ckpt_strategy: str = "mcore"
    """Backend for the async write. ``mcore`` needs no extra deps; megatron-core's own
    default ``nvrx`` requires nvidia-resiliency-ext. Only used when async saves are on."""

    def __post_init__(self):
        # Backfill defaults for any keys the user didn't override so an override dict
        # doesn't have to repeat every default just to set one value.
        if self.transformer_config_kwargs is None:
            self.transformer_config_kwargs = {}
        for k, v in DEFAULT_TRANSFORMER_CONFIG_KWARGS.items():
            self.transformer_config_kwargs.setdefault(k, copy.deepcopy(v))
        if self.optimizer_config_kwargs is None:
            self.optimizer_config_kwargs = {}
        for k, v in DEFAULT_MEGATRON_OPTIMIZER_KWARGS.items():
            self.optimizer_config_kwargs.setdefault(k, copy.deepcopy(v))

@dataclass
class MegatronDDPConfig(BaseConfig):
    grad_reduce_in_fp32: bool = True
    overlap_grad_reduce: bool = False
    overlap_param_gather: bool = False
    average_in_collective: bool = True

@dataclass
class MegatronLoraConfig(BaseConfig):
    lora_type: str = "lora"
    merge_lora: bool = True

@dataclass
class TorchProfilerConfig(BaseConfig):
    """``torch.profiler`` config for policy training steps."""

    enable: bool = False
    ranks: List[int] = field(default_factory=lambda: [0])
    save_path: Optional[str] = None
    """Trace output dir. Required when ``enable=True``; must be a local absolute path."""

    # torch.profiler.schedule
    skip_first: int = 10
    """Steps to skip before scheduling begins."""
    wait: int = 0
    warmup: int = 1
    active: int = 1
    """Number of steps recorded per cycle."""
    repeat: int = 1
    """Number of cycles. 0 means forever."""

    # torch.profiler.profile
    activities: List[str] = field(default_factory=lambda: ["cpu", "cuda"])
    record_shapes: bool = True
    profile_memory: bool = False
    with_stack: bool = True
    with_flops: bool = False
    with_modules: bool = False
    export_type: str = "chrome_trace"
    """``chrome_trace`` or ``stacks``; stacks require ``with_stack=True``."""

    def validate(
        self,
        strategy: Optional[str] = None,
        colocate_all: Optional[bool] = None,
        colocate_policy_ref: Optional[bool] = None,
        fsdp_cpu_offload: Optional[bool] = None,
    ) -> None:
        """Fail fast on invalid or known-incompatible profiler settings."""
        if not self.enable:
            return
        if not self.ranks:
            raise ValueError("`torch_profiler_config.ranks` must be non-empty when profiling is enabled.")
        # Avoid implicit relative paths in Ray runtime working dirs.
        if not self.save_path:
            raise ValueError(
                "`torch_profiler_config.save_path` must be set when profiling is enabled. "
                "Use an absolute local path -- Ray workers run from a /tmp/ray runtime "
                "working dir, so a relative path would write traces there."
            )
        from skyrl.backends.skyrl_train.utils.io.io import is_cloud_path

        if is_cloud_path(self.save_path):
            raise ValueError(
                f"`torch_profiler_config.save_path` must be a local path; got cloud URI "
                f"{self.save_path!r}. torch.profiler cannot write to cloud storage."
            )
        # Empty activities record nothing.
        if not self.activities:
            raise ValueError("`torch_profiler_config.activities` must be non-empty when profiling is enabled.")
        bad_activities = [a for a in self.activities if a.lower() not in TORCH_PROFILER_ACTIVITIES]
        if bad_activities:
            raise ValueError(
                f"invalid `torch_profiler_config.activities` entries {bad_activities}. "
                f"Each must be one of {list(TORCH_PROFILER_ACTIVITIES)}."
            )
        if self.export_type not in TORCH_PROFILER_EXPORT_TYPES:
            raise ValueError(
                f"invalid `torch_profiler_config.export_type`: {self.export_type!r}. "
                f"Must be one of {list(TORCH_PROFILER_EXPORT_TYPES)}."
            )
        if self.export_type == "stacks" and not self.with_stack:
            raise ValueError(
                "`torch_profiler_config.export_type='stacks'` requires `with_stack=true` "
                "(torch.profiler.export_stacks needs stack records)."
            )
        for name in ("skip_first", "wait", "warmup", "repeat"):
            value = getattr(self, name)
            if value < 0:
                raise ValueError(f"`torch_profiler_config.{name}` must be >= 0, got {value}.")
        if self.active < 1:
            raise ValueError(f"`torch_profiler_config.active` must be >= 1, got {self.active}.")

        # FSDP manual CPU offload uses swap_tensors, which conflicts with profiler-held
        # parameter refs during colocated runs.
        if strategy == "fsdp" and fsdp_cpu_offload is False and (colocate_all or colocate_policy_ref):
            raise ValueError(
                "`torch_profiler_config.enable=true` is incompatible with this FSDP configuration: "
                "with the manual CPU-offload path (`policy.fsdp_config.cpu_offload=false`, the default) "
                "under colocation "
                f"(`placement.colocate_all={colocate_all}`, `placement.colocate_policy_ref={colocate_policy_ref}`), "
                "the trainer offloads models to CPU via `torch.utils.swap_tensors` while the profiler holds "
                "references to their parameters, which crashes mid-run with "
                "`RuntimeError: _apply(): Couldn't swap <param>`. "
                "To profile: set `policy.fsdp_config.cpu_offload=true` (FSDP2-native offload, no swap), or "
                "disable colocation (`placement.colocate_all=false` and `placement.colocate_policy_ref=false`), "
                "or use the Megatron backend (`trainer.strategy=megatron`)."
            )