Configuration

@dataclass
class SkyRLTrainConfig(BaseConfig):
    """Root configuration object for SkyRL training with the ``fsdp`` and ``megatron`` backends.

    Every field is overridable from the CLI in ``key.path=value`` form (see
    ``from_cli_overrides``): ``trainer.policy.model.path=...`` sets
    ``SkyRLTrainConfig.trainer.policy.model.path``.
    """

    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}
                  Dict values are serialized as JSON, so ``None``, bools, strings,
                  lists and nested dicts keep their types.

        Returns:
            A fully constructed SkyRLTrainConfig with CLI overrides applied.

        Raises:
            ValueError: If an argument uses the unsupported '+' prefix.
        """
        if isinstance(args, dict):
            args = overrides_dict_to_dotlist(args)

        # 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}
                  Dict values are serialized as JSON, so ``None``, bools, strings,
                  lists and nested dicts keep their types.

        Returns:
            A fully constructed SkyRLTrainConfig with CLI overrides applied.

        Raises:
            ValueError: If an argument uses the unsupported '+' prefix.
        """
        if isinstance(args, dict):
            args = overrides_dict_to_dotlist(args)

        # 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):
    """Training and validation dataset configuration.

    NOTE: datasets are currently loaded entirely into memory, so the maximum usable dataset
    size is bounded by the CPU memory available on a worker node.
    """

    train_data: List[str] = field(default_factory=lambda: [os.path.expanduser("~/data/gsm8k/train.parquet")])
    """Files for the training dataset.
    Each entry is a path to a parquet or json file, or the name of a HuggingFace dataset."""
    val_data: List[str] = field(default_factory=lambda: [os.path.expanduser("~/data/gsm8k/validation.parquet")])
    """Files for the evaluation dataset, in the same formats accepted by ``train_data``.
    When more than one is given, evaluation runs over all of them: both per-dataset metrics
    (keyed by each sample's ``data_source``) and aggregated ``eval/all/*`` metrics are logged,
    and ``trainer.dump_eval_results`` dumps the per-dataset and aggregated results."""
    dataloader: DataLoaderConfig = field(default_factory=DataLoaderConfig)

@dataclass
class DataLoaderConfig(BaseConfig):
    num_workers: Optional[int] = field(
        default=None,
        metadata={
            "help": (
                "Prompt DataLoader worker processes. Default of None auto-derives to 8. "
                "Set 0 for in-process loading that never respawns workers at epoch boundaries."
            )
        },
    )
    persistent_workers: bool = field(
        default=False,
        metadata={
            "help": (
                "Keep DataLoader workers alive across epochs instead of respawning them at "
                "every epoch boundary. Setting this requires `num_workers > 0`"
            )
        },
    )

    def __post_init__(self) -> None:
        if self.num_workers is not None and self.num_workers < 0:
            raise ValueError(f"data.dataloader.num_workers must be None or >= 0, got {self.num_workers}.")

@dataclass
class ModelConfig(BaseConfig):
    path: Optional[str] = None
    """HuggingFace model path (or local directory) for this model."""
    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):
    """LoRA configuration for parameter-efficient fine-tuning.

    Trains a small number of additional low-rank matrices instead of the full model weights.
    """

    rank: int = 0
    """Rank of the low-rank decomposition.
    ``0`` disables LoRA. Higher values increase capacity but also memory usage; 8, 16, 32, and 64 are common choices."""
    alpha: int = 16
    """Scaling factor for LoRA updates."""
    dropout: float = 0.0
    """Dropout probability applied to LoRA layers, to help prevent overfitting."""
    lora_sync_path: str = "/tmp/skyrl_lora_sync"
    """Directory where LoRA adapter weights are saved and synchronized between the training and inference processes.
    Must be accessible to all workers in distributed setups."""
    target_modules: str = "all-linear"
    """Modules to apply LoRA to.
    ``"all-linear"`` targets every linear layer for FSDP/PEFT, and is remapped to a fixed module list
    on Megatron. A list of specific module names can be given instead."""
    exclude_modules: Optional[str] = None
    """Modules to exclude from LoRA."""
    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.

    Reduces GPU memory fragmentation across the offload/backload and forward/backward cycles.
    Automatically turned off around CUDA-IPC weight sync, since IPC handles are incompatible with the
    VMM addresses expandable segments uses; under ``colocate_all=False`` weight sync uses NCCL
    broadcast instead, so it stays on continuously.
    ``InferenceEngineConfig.use_expandable_segments`` is the independent inference-engine knob (the
    trainer and the engine are separate processes with separate allocators). See
    https://docs.skyrl.ai/docs/troubleshooting/troubleshooting for the fragmentation symptoms this
    addresses."""
    sequence_parallel_backend: str = "ulysses"
    strategy: str = "fsdp"
    """Training backend: either ``"fsdp"`` or ``"megatron"``.
    ``"fsdp"`` uses PyTorch's composable ``fully_shard`` API (formerly known as FSDP2)."""
    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
    """Use gradient checkpointing (activation recomputation) to trade compute for memory."""
    gradient_checkpointing_use_reentrant: bool = False
    seed: int = 42
    """Random seed for training."""
    resume_mode: Optional[str] = "latest"
    """``None``/``"none"``, ``"latest"``, or ``"from_path"``.
    See https://docs.skyrl.ai/docs/checkpointing-logging/checkpointing"""
    resume_path: Optional[str] = None
    """Checkpoint directory to resume from. Only used when ``resume_mode="from_path"``."""
    log_path: str = "/tmp/skyrl-logs"
    """Path for infrastructure log files.
    vLLM engine startup, model loading, and worker initialization logs are written to
    ``{log_path}/infra-YYMMDD_HHMMSS.log``. For multi-node training, use a shared filesystem path to
    consolidate logs into a single file. See https://docs.skyrl.ai/docs/checkpointing-logging/logging"""
    ckpt_path: str = field(default_factory=lambda: os.path.expanduser("~/ckpts/"))
    """Directory for resumable training checkpoints (model state, optimizer state, etc.).
    Accepts a local directory path or a cloud storage path (S3, GCS)."""
    max_ckpts_to_keep: int = -1
    """``-1`` to keep all checkpoints, ``N`` to keep only the last N."""
    ckpt_interval: int = 10
    """Save a full training checkpoint every N steps."""
    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.).
    For sharded multi-node HF exports with ``policy.megatron_config.hf_export_config.distributed_save=True``, this must
    be a shared filesystem path visible to all Megatron ranks."""
    bf16: bool = True
    epochs: int = 1
    """Number of epochs (passes over the full dataset)."""
    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.
    Equivalent to the concept of "PPO epochs", where the same experience is iterated over multiple times."""
    train_batch_size: int = 1024
    """Batch size of prompts used for each dataloader step.

    See ``utils/utils.py::validate_batch_sizes`` for the constraints relating this to
    ``policy_mini_batch_size``, ``micro_train_batch_size_per_gpu``, and
    ``micro_forward_batch_size_per_gpu``."""
    policy_mini_batch_size: int = 256
    """Mini batch size for the RL training step; each mini batch is one optimizer step.
    For example, with ``train_batch_size=4`` and ``policy_mini_batch_size=2`` there are 2 optimizer steps (model
    updates) per training batch. This is the *global* mini batch size, counted in prompts — the
    per-worker mini batch is
    ``policy_mini_batch_size * generator.n_samples_per_prompt / number of DP ranks``."""
    critic_mini_batch_size: int = 256
    """Like ``policy_mini_batch_size``, but for the critic model.
    The critic generally tolerates off-policy updates better than the policy, so setting this lower than
    ``policy_mini_batch_size`` (i.e. more critic updates) is usually preferable."""
    micro_train_batch_size_per_gpu: int = 1
    """Micro batch size during the training step, common to both policy and critic.
    Each mini batch is split into micro batches of this size, and gradients are accumulated over them."""
    micro_forward_batch_size_per_gpu: int = 1
    """Micro batch size during the forward pass, i.e. log probability or value computation.
    Common to both policy and critic. Each mini batch is split into micro batches of this size."""
    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
    """Re-sync the reference model from the policy model at every epoch boundary."""
    remove_microbatch_padding: bool = True
    """Pack samples into the THD layout and strip intra-microbatch padding (requires flash attention).
    Common to all models."""
    eval_batch_size: int = 1024
    """Batch size for evaluation."""
    eval_before_train: bool = True
    """Evaluate the model once before training starts."""
    eval_interval: int = 5
    """Evaluate against the validation dataset every N steps. ``-1`` to disable evaluation."""
    max_prompt_length: int = 512
    """Maximum prompt length during training.
    Prompts longer than this are filtered out of the train/eval datasets at load time, not
    truncated."""
    flash_attn: bool = True
    disable_fast_tokenizer: bool = False
    project_name: str = "skyrl"
    """Project name in WandB and MLflow."""
    run_name: str = "test_run"
    """Run name in WandB and MLflow."""
    logger: str = "wandb"
    """Logger to use: ``"wandb"``, ``"mlflow"``, ``"swanlab"``, ``"tensorboard"``, or ``"console"``.
    See https://docs.skyrl.ai/docs/checkpointing-logging/logging"""
    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 each training data batch to a file for debugging.
    The batch at global step N is written to
    ``{export_path}/dumped_data/global_step_{N}_training_input.pkl``."""
    dump_eval_results: bool = True
    """Dump full evaluation results to a file.
    Results at global step N are written to ``{export_path}/dumped_evals/global_step_{N}_evals``, with both per-dataset
    and aggregated results when multiple validation datasets are configured."""
    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):
    """Optimizer configuration, shared by the policy and critic models."""

    lr: float = 1e-6
    """Learning rate."""
    adam_betas: List[float] = field(default_factory=lambda: [0.9, 0.999])
    """Betas for the AdamW optimizer."""
    weight_decay: float = 1e-2
    """L2 regularization strength for AdamW."""
    max_grad_norm: float = 1.0
    """Gradient clipping. The total L2 norm of the model gradients is scaled to this value."""
    offload_after_step: bool = True
    """Offload optimizer state to CPU after each full training step.
    Applies under colocation (``colocate_all``, or ``colocate_policy_ref`` for policy/ref), and is
    inert when ``fsdp_config.cpu_offload=True`` since FSDP2 then offloads natively. Without
    colocation it can be preferable to leave optimizer state on GPU, avoiding both the offload cost
    and the extra CPU memory usage."""
    num_warmup_steps: int = 0
    """Number of mini-batch steps to warmup the optimizer."""
    scheduler: str = "constant_with_warmup"
    """Learning rate scheduler. Intended to align with ``transformers.SchedulerType``:
    https://huggingface.co/docs/transformers/main/en/main_classes/optimizer_schedules#transformers.SchedulerType"""

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

@dataclass
class FakeInt4QatConfig(BaseConfig):
    """Fake-INT4 quantization-aware training for MoE experts (Megatron only).

    When the inference engine serves the MoE experts as real ``compressed-tensors``
    INT4 (e.g. ``casperhansen/Qwen3.6-35B-A3B-INT4-RTN``) but the trainer holds
    BF16 masters, enabling this fake-quantizes the frozen expert GEMMs onto the
    same INT4 grid in the forward pass (straight-through backward), removing the
    train/infer weight mismatch. See
    ``skyrl.backends.skyrl_train.workers.megatron.fake_int4_qat``.
    """

    enabled: bool = False
    group_size: int = 32
    """Group size along the input dim; must match the served checkpoint (32)."""
    symmetric: bool = True
    scale_divisor: float = 7.5
    """Symmetric-INT4 scale divisor ``scale = amax / scale_divisor``:
    ``7.5`` = llm-compressor / compressed-tensors RTN (``[-8, 7]``; matches
    ``casperhansen/Qwen3.6-35B-A3B-INT4-RTN``); ``7.0`` = Kimi K2-Thinking / K2.6 /
    Miles (``[-7, 7]``). Set ``q_min`` consistently."""
    q_min: float = -8.0
    """Lower clamp of the INT4 code range: ``-8`` for llm-compressor RTN
    (``scale_divisor=7.5``), ``-7`` for Kimi/Miles (``scale_divisor=7.0``, whose
    QAT never emits ``-8``)."""
    bf16_base_path: Optional[str] = None
    """Megatron-Bridge cannot load a compressed-tensors INT4 checkpoint, so when
    ``model.path`` points at the INT4 model the trainer loads its BF16 master
    weights from this path instead. The INT4 ``model.path`` remains what the
    inference engine serves and the logical name. When None, the trainer loads
    weights from ``model.path`` directly (only valid if that path is already a
    BF16 checkpoint)."""

@dataclass
class MTPConfig(BaseConfig):
    enabled: bool = False
    """Whether to train MTP draft heads and use them for speculative decoding."""
    num_speculative_tokens: int = 1
    """Draft depth vLLM speculates per step, independent of the trained head count. Single-head
    checkpoints (Qwen3.5/Qwen3-Next/DeepSeek-V3) reuse the one head autoregressively at depths > 1;
    expect acceptance to decay with depth (the head is trained at depth 1)."""
    loss_weight: float = 0.1
    """Weight ``w`` of the draft loss in ``policy_loss + w * draft_loss``."""

@dataclass
class FSDPConfig(BaseConfig):
    cpu_offload: bool = False
    """Offload params and optimizer state to CPU during the forward pass.

    Corresponds to FSDP2's ``offload_policy``:
    https://docs.pytorch.org/docs/stable/distributed.fsdp.fully_shard.html
    Enabling this replaces SkyRL's manual colocation offload (and makes
    ``optimizer_config.offload_after_step`` inert) rather than stacking with it; see
    https://docs.skyrl.ai/docs/tutorials/placement for the difference."""
    reshard_after_forward: Union[bool, int] = True
    """FSDP2 only.
    Accepts True, False, or an int between 1 and ``fsdp_size``. See
    https://docs.pytorch.org/docs/stable/distributed.fsdp.fully_shard.html. Setting False retains the full model
    parameters on each worker (similar to DeepSpeed ZeRO stage 2)."""
    fsdp_size: int = -1
    """Group size within which worker state is sharded, for hybrid sharding in multi-node runs.
    ``-1`` shards across all workers in the group. Example: with 8 workers across 2 nodes (4 each)
    and ``fsdp_size=4``, training state is fully sharded across the 4 ranks within each node and
    replicated (data-parallel) across nodes."""
    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):
    """Megatron-Core backend configuration, used when ``trainer.strategy="megatron"``.

    The parallelism sizes must satisfy:

    - ``model_size = pipeline_model_parallel_size * tensor_model_parallel_size * context_parallel_size``
    - ``dp_size = world_size / model_size``
    - ``world_size % (pipeline_model_parallel_size * expert_model_parallel_size * expert_tensor_parallel_size) == 0``

    The last rule means ``expert_model_parallel_size * expert_tensor_parallel_size`` can scale
    independently of ``tensor_model_parallel_size * context_parallel_size``, and can span data
    parallel ranks. See https://docs.skyrl.ai/docs/examples/megatron for sizing guidance.
    """

    tensor_model_parallel_size: int = 1
    """Tensor model parallel size, reducing memory for model parameters and activations.
    Megatron sequence parallelism (unrelated to Ulysses ``sequence_parallel_size``) is enabled automatically whenever
    this is greater than 1."""
    pipeline_model_parallel_size: int = 1
    """Pipeline model parallel size, sharding model layers across GPUs."""
    context_parallel_size: int = 1
    """Context parallel size, reducing activation memory along the sequence-length dimension."""
    expert_model_parallel_size: int = 1
    """Expert parallel size, sharding expert modules across GPUs."""
    expert_tensor_parallel_size: Optional[int] = None
    """Tensor parallel size for each expert module.
    ``None`` lets Megatron resolve it to ``tensor_model_parallel_size``. Setting this to ``1`` is recommended for best
    performance when ``expert_model_parallel_size > 1``."""
    # 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
    """Enable Megatron router replay.
    Used together with ``generator.inference_engine.enable_return_routed_experts`` to enable R3."""
    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)
    """Pass-through config for Megatron's ``DistributedDataParallelConfig``:
    https://github.com/NVIDIA/Megatron-LM/blob/core_r0.13.0/megatron/core/distributed/distributed_data_parallel_config.py"""
    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)
    )
    """Pass-through kwargs for Megatron's ``OptimizerConfig``.

    https://github.com/NVIDIA/Megatron-LM/blob/core_r0.13.0/megatron/core/optimizer/optimizer_config.py
    Any keys overlapping with what SkyRL resolves from ``optimizer_config`` are overridden by the
    values here. ``*_dtype`` keys accept case-insensitive dtype names, coerced by
    ``distributed.megatron.optimizer_dtype``; see
    https://docs.skyrl.ai/docs/examples/megatron for the accepted names and per-field checks.
    ``use_precision_aware_optimizer=True`` can cause checkpointing to fail
    (https://github.com/nvidia/megatron-lm/issues/1820); leaving it ``False`` is recommended."""
    transformer_config_kwargs: Dict[str, Any] = field(
        default_factory=lambda: copy.deepcopy(DEFAULT_TRANSFORMER_CONFIG_KWARGS)
    )
    """Pass-through kwargs for Megatron's ``TransformerConfig``:
    https://github.com/NVIDIA/Megatron-LM/blob/core_r0.13.0/megatron/core/transformer/transformer_config.py
    Also the place to put HuggingFace config overrides (e.g. ``rope_parameters``) for the Megatron
    backend, where FSDP would use ``model_config_kwargs``."""
    empty_cuda_cache: Optional[bool] = True
    """Manually empty torch's CUDA cache between the forward/backward pass and the optimizer step.
    This frees reserved-but-unallocated memory and can help avoid OOMs in the optimizer."""
    model_config_kwargs: dict = field(default_factory=dict)
    """HF-config overrides read from the nested ``model_config`` key only.
    Used for bridge and RoPE resolution. Not a general HF-config override -- use
    ``transformer_config_kwargs`` instead."""
    dist_ckpt_optim_fully_reshardable: bool = False
    """When True, use the "fully-reshardable" format for the distributed-optimizer checkpoint.
    When False (default), use the "dp-reshardable" format, which is more efficient but only
    supports resharding along the data-parallel dimension. See
    https://github.com/NVIDIA/Megatron-LM/blob/core_v0.16.0/megatron/core/optimizer/distrib_optimizer.py"""
    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."""
    async_save_prestage_to_cpu: bool = False
    """Copy shards to host memory on the training rank before handing them to the async
    checkpoint writer, instead of letting the writer pull them over CUDA IPC.

    Enable this when using async save on machines with restricted ptrace permissions. 
    megatron-core hands the writer the *GPU* tensors and copies them in the writer process;
    with ``expandable_segments:True`` those handles are file descriptors that the writer 
    imports via ``pidfd_getfd``, which needs ptrace-attach permission on the rank. 
    Where that is refused (e.g. ``kernel.yama.ptrace_scope=1``, as on some CI runners) 
    the writer dies with ``pidfd_getfd: Operation not permitted`` and every rank then 
    hangs on the preload barrier.

    Off by default: prestaging shrinks the rank's blocking
    window but delays the background write (per-tensor shared-memory handoff cost).
    See ``_stage_async_request_to_host``."""

    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))