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

        # generator rope params default to trainer rope params
        if self.generator.rope_scaling is None and self.trainer.rope_scaling is not None:
            self.generator.rope_scaling = self.trainer.rope_scaling
        if self.generator.rope_theta is None and self.trainer.rope_theta is not None:
            self.generator.rope_theta = self.trainer.rope_theta
        # 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:
            # TODO(Charlie): debug why inference http endpoint is slow when num_workers is 8
            self.data.dataloader.num_workers = 0 if self.generator.inference_engine.enable_http_endpoint else 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 either"
                " set explicitly to 0 or forced to 0 by the inference HTTP endpoint."
            )

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

@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)
    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)
    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
    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
    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
    rope_scaling: Optional[Dict[str, Any]] = None
    rope_theta: Optional[float] = None
    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."""

    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.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}."
            )

@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)
    torch_profiler_config: MegatronTorchProfilerConfig = field(default_factory=MegatronTorchProfilerConfig)
    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."""

    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 MegatronTorchProfilerConfig(BaseConfig):
    enable: bool = False
    ranks: List[int] = field(default_factory=list)
    save_path: Optional[str] = None

@dataclass
class PlacementConfig(BaseConfig):
    colocate_all: bool = True
    """When True, training and inference share the same GPUs."""
    colocate_policy_ref: bool = True
    """When colocate_all is False, True (default) still colocates policy and ref
    on the same GPUs (one shared placement group). Set this item to False to place
    policy and ref on separate GPUs (their own placement groups); needed when
    a large model's policy and ref shards can't both fit on one GPU."""
    policy_num_nodes: int = 1
    policy_num_gpus_per_node: int = 1
    critic_num_nodes: int = 1
    critic_num_gpus_per_node: int = 1
    ref_num_nodes: int = 1
    ref_num_gpus_per_node: int = 1

@dataclass
class PolicyConfig(BaseConfig):
    model: ModelConfig = field(default_factory=lambda: copy.deepcopy(ModelConfig(path="Qwen/Qwen2.5-1.5B-Instruct")))
    optimizer_config: OptimizerConfig = field(default_factory=OptimizerConfig)
    fsdp_config: FSDPConfig = field(default_factory=FSDPConfig)
    sequence_parallel_size: int = 1
    use_torch_compile: bool = False
    """Apply torch.compile to logits calculation."""
    record_memory: bool = False
    """Save memory snapshots to ``{ckpt_path}/memory_snapshots/``.
    Visualize by dragging pickle files to https://docs.pytorch.org/memory_viz."""
    megatron_config: MegatronConfig = field(default_factory=MegatronConfig)
    model_config_kwargs: dict = field(default_factory=dict)
    """Pass-through kwargs for the HuggingFace model config (FSDP backends).
    For Megatron, use ``policy.megatron_config.transformer_config_kwargs`` instead."""
    language_model_only: bool = False
    """When True, skip vision encoder initialization for multimodal models (e.g. Qwen3.5).
    Loads only the language model backbone using AutoModelForCausalLM."""
    inference_only_init: bool = False
    """When True, set up the policy worker for inference-only flows (forward + weight
    sync, no train_step), skipping the training-only state that would otherwise OOM
    memory-constrained nodes (e.g. large MoE on 4xH100). NOT valid for actual training.
    Backend-specific behavior:
    - FSDP: initialize weights in bf16 instead of fp32 (skipping the fp32 master weights
      that mixed-precision training requires) and skip optimizer/LR-scheduler construction.
    - Megatron: skip optimizer/LR-scheduler construction (DistributedOptimizer eagerly
      materializes fp32 master + AdamW state on GPU)."""

@dataclass
class CriticConfig(BaseConfig):
    model: ModelConfig = field(default_factory=ModelConfig)
    optimizer_config: OptimizerConfig = field(default_factory=lambda: OptimizerConfig(lr=5e-6))
    fsdp_config: FSDPConfig = field(default_factory=FSDPConfig)
    sequence_parallel_size: int = 1
    model_config_kwargs: dict = field(default_factory=dict)

@dataclass
class RefConfig(BaseConfig):
    model: ModelConfig = field(default_factory=ModelConfig)
    sequence_parallel_size: int = 1
    fsdp_config: FSDPConfig = field(default_factory=FSDPConfig)
    megatron_config: MegatronConfig = field(default_factory=MegatronConfig)
    model_config_kwargs: dict = field(default_factory=dict)
    language_model_only: bool = False
    """When True, skip vision encoder initialization for multimodal models (e.g. Qwen3.5).
    Loads only the language model backbone using AutoModelForCausalLM."""