cli_args.py

import argparse
import json
import os
from dataclasses import MISSING as dataclass_missing
from dataclasses import asdict, dataclass, field, fields
from pathlib import Path
from typing import TYPE_CHECKING, Any, ClassVar, TypeVar

import uvloop
import yaml
from omegaconf import MISSING, DictConfig, OmegaConf

from astraflow.train_worker.utils import logging, name_resolve, pkg_version
from astraflow.train_worker.utils.pkg_version import is_version_less

if TYPE_CHECKING:
    from transformers import PreTrainedTokenizerFast

uvloop.install()

logger = logging.getLogger("CLI args")

ConfigT = TypeVar("ConfigT")


@dataclass
class NormConfig:
    """Configuration for reward/advantage normalization."""

    mean_level: str | None = field(
        default="batch",
        metadata={
            "help": "Mean level for normalization. None for no mean normalization.",
            "choices": ["batch", "group", None],
        },
    )
    mean_leave1out: bool = field(
        default=False,
        metadata={"help": "Whether to use leave-one-out average."},
    )
    std_level: str | None = field(
        default="batch",
        metadata={
            "help": "Standard deviation level for normalization. None for no std normalization.",
            "choices": ["batch", "group", None],
        },
    )
    std_unbiased: bool = field(
        default=True,
        metadata={
            "help": "Whether to use unbiased standard deviation computation. Defaults to True (changed from False in v0.3.4)."
        },
    )
    eps: float = field(
        default=1e-5,
        metadata={
            "help": "The eps when dividing by standard deviation to avoid numerical issues."
        },
    )
    group_size: int | None = field(
        default=None, metadata={"help": "Group size for group-level normalization. Auto-set from n_samples if omitted."}
    )


@dataclass
class MicroBatchSpec:
    """Specification for splitting micro-batches during training."""

    n_mbs: int | None = field(
        default=1,
        metadata={
            "help": "Number of micro-batches (or minimum number if max_tokens_per_mb is set). Used when max_tokens_per_mb is None or as minimum count",
        },
    )
    granularity: int = field(
        default=1,
        metadata={
            "help": "Granularity of each micro-batch. Adjacent sequences are grouped by this size when dividing microbatches.",
        },
    )
    max_tokens_per_mb: int | None = field(
        default=None,
        metadata={
            "help": "Maximum tokens per micro-batch for each forward pass. When set, n_mbs becomes the minimum number of micro-batches.",
        },
    )
    n_mbs_divisor: int = field(
        default=1,
        metadata={
            "help": "Divisor for the number of micro-batches. The final number of micro-batches will be adjusted to be divisible by this value.",
        },
    )

    @classmethod
    def new(cls, mb_spec: "MicroBatchSpec", **kwargs):
        """Create new spec with updated fields while maintaining Omegaconf compatibility."""
        fields = dict(
            n_mbs=mb_spec.n_mbs,
            granularity=mb_spec.granularity,
            max_tokens_per_mb=mb_spec.max_tokens_per_mb,
            n_mbs_divisor=mb_spec.n_mbs_divisor,
        )
        fields.update(kwargs)
        return cls(**fields)


@dataclass
class GenerationHyperparameters:
    """Controls text generation behavior for rollout."""

    n_samples: int = field(
        default=1, metadata={"help": "Number of sequences to generate per prompt."}
    )
    max_new_tokens: int = field(
        default=16384, metadata={"help": "Maximum number of tokens to generate."}
    )
    min_new_tokens: int = field(
        default=0, metadata={"help": "Minimum number of tokens to generate."}
    )
    max_tokens: int = field(
        default=65536,
        metadata={
            "help": "Maximum number of tokens including prompt and generated tokens."
        },
    )
    greedy: bool = field(
        default=False,
        metadata={"help": "Whether to use greedy decoding (max probability)."},
    )
    top_p: float = field(
        default=1.0,
        metadata={"help": "Nucleus sampling probability threshold (0.0, 1.0]."},
    )
    top_k: int = field(
        default=int(1e8),
        metadata={"help": "Number of highest probability tokens to consider."},
    )
    temperature: float = field(
        default=1.0,
        metadata={"help": "Sampling temperature. Higher values increase diversity."},
    )
    stop_token_ids: list[int] = field(
        default_factory=list,
        metadata={"help": "Stop generation when encountering these token IDs."},
    )
    ignore_eos: bool = field(
        default=False,
        metadata={"help": "Do not stop generation when EOS is encountered."},
    )
    skip_special_tokens: bool = field(
        default=True,
        metadata={"help": "Skip special tokens when decoding/displaying outputs."},
    )
    include_pad_in_stop_tokens: bool = field(
        default=True,
        metadata={
            "help": "Whether to include PAD token in stop_token_ids (EOS is always included)."
        },
    )
    stop: list[str] | None = field(
        default=None,
        metadata={
            "help": "One or multiple stop words. Generation will stop if one of these words is sampled."
        },
    )
    frequency_penalty: float = field(
        default=0.0,
        metadata={
            "help": (
                "Penalizes tokens based on their frequency in generation so far. "
                "Must be between -2 and 2 where negative numbers encourage repetition."
            )
        },
    )
    lora_name: str = field(
        default="",
        metadata={"help": "Lora name to be used for this generation."},
    )
    use_beam_search: bool = field(
        default=False,
        metadata={
            "help": "Enable beam search in the vLLM engine. When enabled, sampling parameters like temperature, top-p, and top-k are auto ignored."
        },
    )
    # NOTE: to add new parameters, please correctly handle them in the `to_openai_args_dict` method.

    def new(self, **kwargs):
        args = asdict(self)
        args.update(kwargs)
        return GenerationHyperparameters(**args)

    def new_with_stop_and_pad_token_ids(self, tokenizer: "PreTrainedTokenizerFast"):
        """Create a new generation hyperparameters with stop and pad token ids added."""
        new_stop_token_ids = self.stop_token_ids.copy()
        if (
            self.include_pad_in_stop_tokens
            and tokenizer.pad_token_id is not None
            and tokenizer.pad_token_id not in new_stop_token_ids
        ):
            new_stop_token_ids.append(tokenizer.pad_token_id)
            print(f"Appended pad token id: {tokenizer.pad_token_id} to stop_token_ids")
        if (
            tokenizer.eos_token_id is not None
            and tokenizer.eos_token_id not in new_stop_token_ids
        ):
            new_stop_token_ids.append(tokenizer.eos_token_id)
            print(f"Appended eos token id: {tokenizer.eos_token_id} to stop_token_ids")
        return self.new(stop_token_ids=new_stop_token_ids)

    def to_openai_completions_args_dict(
        self, exclude_args: list[str] | None = None
    ) -> dict[str, Any]:
        return self.to_openai_args_dict(
            exclude_args=exclude_args, api_format="completions"
        )

    def to_openai_responses_args_dict(
        self, exclude_args: list[str] | None = None
    ) -> dict[str, Any]:
        return self.to_openai_args_dict(
            exclude_args=exclude_args, api_format="responses"
        )

    def to_openai_agents_model_settings_dict(
        self, exclude_args: list[str] | None = None
    ) -> dict[str, Any]:
        return self.to_openai_args_dict(
            exclude_args=exclude_args, api_format="openai-agents"
        )

    _OPENAI_UNSUPPORTED_ARGS: ClassVar[set[str]] = {
        "min_new_tokens",  # Not supported by OpenAI
        "greedy",  # Not directly supported by OpenAI
        "top_k",  # Not supported by OpenAI
        "stop_token_ids",  # Not supported by OpenAI
        "ignore_eos",  # Not supported by OpenAI
        "skip_special_tokens",  # Not supported by OpenAI
        "lora_name",  # Not supported by OpenAI
        "use_beam_search",  # Not supported by OpenAI
        "max_tokens",  # deprecated by "completions", not used in "responses", should be `max_new_tokens` in "openai-agents"
    }

    def to_openai_args_dict(
        self, exclude_args: list[str] | None = None, api_format: str = "completions"
    ) -> dict[str, Any]:
        """Convert the generation hyperparameters to a dictionary of arguments for OpenAI client."""
        final_exclude_args = set(exclude_args) if exclude_args is not None else set()
        final_exclude_args.update(self._OPENAI_UNSUPPORTED_ARGS)
        # TODO: move the excluded args into extra body, so they can be passed through the client request

        mapping = {"n_samples": "n"}
        if api_format == "completions":
            mapping["max_new_tokens"] = "max_completion_tokens"
        elif api_format == "responses":
            mapping["max_new_tokens"] = "max_output_tokens"
        elif api_format == "openai-agents":
            # NOTE: max_tokens in openai-agents means `max_new_tokens` in sglang/vllm. This is not a bug
            mapping["max_new_tokens"] = "max_tokens"
        else:
            raise ValueError(f"Unsupported API format: {api_format}")

        res = {}
        for k, v in asdict(self).items():
            if k in final_exclude_args:
                should_warn = False

                current_value = getattr(self, k)
                f = next(_field for _field in fields(self) if _field.name == k)

                # Check if equal to the default value
                if f.default is not dataclass_missing:
                    if current_value != f.default:
                        should_warn = True
                elif f.default_factory is not dataclass_missing:
                    if current_value != f.default_factory():
                        should_warn = True

                if should_warn:
                    logger.warning(
                        f"Unsupported arg for openai format: `{k}` with value {current_value}"
                    )
                continue
            key = mapping.get(k, k)
            if key in res:
                logger.warning(f"Overriding key: {key} from {k} with value: {v}")
            res[key] = v

        return res


# Train Engine Configs


@dataclass
class OptimizerConfig:
    """Configuration for model optimization during training."""

    type: str = field(
        default="adam",
        metadata={
            "help": "Optimizer type. Adam_bf16 currently only supported FSDP Engine.",
            "choices": ["adam", "sgd", "adam_bf16"],
        },
    )
    lr: float = field(default=1e-3, metadata={"help": "Learning rate"})
    weight_decay: float = field(default=0.01, metadata={"help": "Weight decay"})
    beta1: float = field(
        default=0.9,
        metadata={
            "help": "Adam beta1 parameter. Only effective when optimizer_type is adam/adam_bf16"
        },
    )
    beta2: float = field(
        default=0.999,
        metadata={
            "help": "Adam beta2 parameter. Only effective when optimizer_type is adam/adam_bf16"
        },
    )
    eps: float = field(
        default=1e-8,
        metadata={
            "help": "Adam epsilon parameter. Only effective when optimizer_type is adam/adam_bf16"
        },
    )
    min_lr_ratio: float = field(
        default=0.0,
        metadata={
            "help": "Minimum learning rate ratio after annealing",
        },
    )
    lr_scheduler_type: str = field(
        default="constant",
        metadata={
            "help": "Learning rate scheduler type",
            "choices": ["linear", "cosine", "constant"],
        },
    )
    warmup_steps_proportion: float = field(
        default=0.0,
        metadata={
            "help": "Proportion of training steps for warmup",
        },
    )
    offload: bool = field(
        default=False, metadata={"help": "Enable optimizer state offloading"}
    )
    initial_loss_scale: float = field(
        default=2**32, metadata={"help": "Initial loss scaling factor"}
    )
    min_loss_scale: float = field(
        default=1.0, metadata={"help": "Minimum loss scaling factor"}
    )
    loss_scale_window: float = field(
        default=5, metadata={"help": "Window size for loss scaling adjustment"}
    )
    hysteresis: int = field(
        default=2, metadata={"help": "Hysteresis (scaling factor) for loss scaling"}
    )
    gradient_clipping: float = field(
        default=1.0, metadata={"help": "Gradient clipping threshold"}
    )


@dataclass
class FSDPWrapPolicy:
    """Policy configuration for FSDP model layer wrapping. None defaults to wrapping transformer decoder layers defined by transformers."""

    transformer_layer_cls_to_wrap: list[str] | None = field(
        default=None,
        metadata={"help": "A list of transformer layer names for FSDP to wrap."},
    )


@dataclass
class FSDPEngineConfig:
    """Configuration for Fully Sharded Data Parallel (FSDP) training backend."""

    wrap_policy: FSDPWrapPolicy | None = field(
        default=None,
        metadata={"help": "FSDP wrap policy, specifying model layers to wrap."},
    )
    offload_params: bool = field(
        default=False,
        metadata={"help": "Whether to offload FSDP parameters to CPU."},
    )


# These configurations are used by Megatron Bridge to build Megatron models.
@dataclass
class DistributedDataParallelConfig:
    """Configuration for Megatron's DistributedDataParallel.
    Refer to Megatron-LM documentation for details.
    """

    grad_reduce_in_fp32: bool = True
    overlap_grad_reduce: bool = False
    overlap_param_gather: bool = False
    align_param_gather: bool = False
    use_distributed_optimizer: bool = True
    check_for_nan_in_grad: bool = False
    bucket_size: int | None = None
    average_in_collective: bool = False
    fp8_param_gather: bool = False


@dataclass
class MegatronEngineConfig:
    """Configuration for Megatron-LM training framework.
    Refer to Megatron-LM documentation for implementation details.
    """

    # Distributed Training Configuration
    wrap_with_ddp: bool = True
    use_torch_fsdp2: bool = False  # TODO: pending test
    use_custom_fsdp: bool = False  # TODO: pending test
    ddp: DistributedDataParallelConfig = field(
        default_factory=DistributedDataParallelConfig
    )
    virtual_pipeline_parallel_size: int = field(
        default=1,
        metadata={
            "help": (
                "Virtual pipeline parallel size for Megatron interleaved schedule. "
                "Set to >1 to enable VPP. Default is 1 (disabled)."
            )
        },
    )
    # Don't use MegatronOptimizerConfig here because OmegaConf
    # does not recognize the annotation "torch.dtype"
    overlap_param_gather_with_optimizer_step: bool = False

    # Precision Configuration
    use_precision_aware_optimizer: bool = False
    main_grads_dtype: str = "float32"
    main_params_dtype: str = "float32"
    exp_avg_dtype: str = "float32"
    exp_avg_sq_dtype: str = "float32"

    # Checkpointing Configuration
    async_save: bool = False
    use_checkpoint_opt_param_scheduler: bool = True

    # Deterministic Option
    # NOTE: This option forces torch to use deterministic algorithms,
    # which makes sure that two forward passes with the same input
    # will produce the same output. However, it may have a performance impact.
    # It is recommended to set this option to True for RL training on MoE models for stability.
    use_deterministic_algorithms: bool = False

    # Gradient checkpointing options, only effective when gradient_checkpointing=True
    recompute_granularity: str | None = "full"
    recompute_method: str | None = "uniform"
    recompute_num_layers: int | None = 1
    distribute_saved_activations: bool | None = None
    recompute_modules: list[str] | None = None


@dataclass
class SchedulingStrategy:
    type: str = field(
        default="separation", metadata={"choices": ["separation", "colocation"]}
    )
    target: str | None = field(
        default=None, metadata={"help": "The target role to be colocated with"}
    )


@dataclass
class TrainEngineConfig:
    """Core configuration for model training, including optimization and backend settings."""

    experiment_name: str = ""
    trial_name: str = ""
    path: str = field(default="", metadata={"help": "Path to HuggingFace checkpoint"})
    attn_impl: str = field(
        default="kernels-community/flash-attn2",
        metadata={
            "help": (
                "Attention implementation for huggingface transformers model. "
                "Default pulls a prebuilt FlashAttention-2 kernel from the HF kernels "
                "hub (ABI-matched to torch, incl. varlen for packed sequences). The "
                "literal 'flash_attention_2' loads the local flash-attn wheel, which is "
                "ABI-broken on torch>=2.11+cu13; 'sdpa' works but relies on position_ids "
                "resets for packed block-diagonal masking."
            ),
            "choices": [
                "kernels-community/flash-attn2",
                "flash_attention_2",
                "sdpa",
                "eager",
            ],
        },
    )
    init_from_scratch: bool = field(
        default=False, metadata={"help": "Initialize model weights randomly"}
    )
    is_critic: bool = field(
        default=False,
        metadata={"help": "Whether to use a critic/reward model"},
    )
    temperature: float = field(
        default=1.0, metadata={"help": "Temperature during generation."}
    )
    # Runtime microbatch limit
    mb_spec: MicroBatchSpec = field(default_factory=MicroBatchSpec)
    pad_to_maximum: bool = field(
        default=False,
        metadata={
            "help": (
                "Whether to pad each microbatch to the length upper bound specified by mb_spec. "
                "Can reduce memory fragmentation but slows down training."
            )
        },
    )

    # Training Backend Configuration
    disable_dropout: bool = field(
        default=False, metadata={"help": "Disable dropout layers during training"}
    )
    gradient_checkpointing: bool = field(
        default=False, metadata={"help": "Enable gradient checkpointing"}
    )
    dtype: str = field(default="bfloat16", metadata={"help": "Parameter data type."})
    grad_reduce_dtype: str = field(
        default="float32", metadata={"help": "Gradient reduction data type."}
    )
    optimizer: OptimizerConfig | None = field(
        default=None,
        metadata={"help": "Optimizer configuration. None means no training."},
    )

    fsdp: FSDPEngineConfig = field(default_factory=FSDPEngineConfig)
    megatron: MegatronEngineConfig = field(default_factory=MegatronEngineConfig)

    # Lora
    use_lora: bool = field(
        default=False,
        metadata={
            "help": "Whether to use LoRA. Only support FSDP. Note that should be enabled together with vLLM/SGLang."
        },
    )
    lora_rank: int = field(default=32, metadata={"help": "lora rank"})
    lora_alpha: int = field(default=16, metadata={"help": "lora alpha"})
    target_modules: list[str] = field(
        default_factory=list,
        metadata={"help": "lora target_modules."},
    )
    peft_type: str = field(
        default="lora",
        metadata={"help": "peft method type. Only LoRA is supported for now."},
    )
    scheduling_strategy: SchedulingStrategy = field(
        default_factory=SchedulingStrategy,
        metadata={
            "help": "The scheduling strategy of this TrainEngine, either separation or colocation."
        },
    )


@dataclass
class PPOActorConfig(TrainEngineConfig):
    """Configuration for PPO actor model, a subclass of a TrainEngine."""

    # Core PPO/GRPO Parameters
    group_size: int | None = field(
        default=None, metadata={"help": "Number of sequences in each group. Auto-set from n_samples if omitted."}
    )
    temperature: float | None = field(
        default=None, metadata={"help": "Temperature for importance weight calculation. Auto-set from top-level temperature if omitted."}
    )
    ppo_n_minibatches: int = field(
        default=4, metadata={"help": "Number of minibatches for each PPO update"}
    )
    filter_zero_adv_in_batch: bool = field(
        default=True,
        metadata={
            "help": "If True, drop samples whose group_reward_std == 0 from each "
            "fetched training batch on rank 0 before DP dispatch. Per-step batch "
            "size becomes variable. Independent of buffer.filter_function — both "
            "can be enabled."
        },
    )
    eps_clip: float = field(
        default=0.2, metadata={"help": "Clipping factor for policy ratio"}
    )
    eps_clip_higher: float | None = field(
        default=None,
        metadata={
            "help": "Clipping factor (higher value) for policy ratio. Default is None. When eps_clip_higher is set (decoupled), eps_clip will be used as the lower value."
        },
    )
    c_clip: float | None = field(
        default=None,
        metadata={
            "help": "Dual clipping factor for policy ratio, must be > 1.0. None disables dual clipping."
        },
    )
    # M2PO
    m2_threshold: float | None = field(
        default=None, metadata={"help": "The second momentum threshold for M2PO."}
    )
    # Reward
    reward_norm: NormConfig | None = field(
        default=None,
        metadata={"help": "Normalization configuration for rewards"},
    )
    reward_scaling: float = field(
        default=1.0, metadata={"help": "Reward scaling factor"}
    )
    reward_bias: float = field(default=0.0, metadata={"help": "Reward bias"})
    reward_clip: float = field(
        default=20.0, metadata={"help": "Maximum absolute value for reward clipping"}
    )
    overlong_reward_penalty: bool = field(
        default=False,
        metadata={"help": "Penalty for overlong sequences. Used within DAPO."},
    )
    overlong_tokens: int | None = field(
        default=None,
        metadata={"help": "Number of tokens in the tail that will receive a penalty"},
    )
    overlong_penalty_factor: float | None = field(
        default=None,
        metadata={"help": "Penalty factor for tokens in the tail"},
    )
    mask_no_eos_with_zero: bool = field(
        default=False,
        metadata={
            "help": "Mask truncated generations (no EOS token) and exclude from training"
        },
    )

    # Advantage Estimation
    discount: float = field(
        default=1.0, metadata={"help": "Discount factor for future rewards"}
    )
    gae_lambda: float = field(
        default=1.0, metadata={"help": "Lambda parameter for GAE"}
    )
    adv_norm: NormConfig | None = field(
        default=None, metadata={"help": "Normalization configuration for advantages."}
    )

    # KL Control
    kl_ctl: float = field(default=0.1, metadata={"help": "KL divergence coefficient"})
    kl_penalty_coef: float = field(
        default=0.0,
        metadata={
            "help": "Independent KL penalty coefficient between current policy and reference policy"
        },
    )
    kl_estimator: str = field(
        default="k3",
        metadata={"help": "KL divergence estimator", "choices": ["k1", "k2", "k3"]},
    )

    # SAPO (Soft Adaptive Policy Optimization) - https://arxiv.org/abs/2511.20347
    use_sapo_loss: bool = field(
        default=False,
        metadata={"help": "Use SAPO loss (mutually exclusive with PPO clipping)"},
    )
    sapo_tau_pos: float = field(
        default=1.0,
        metadata={"help": "SAPO temperature for positive advantages"},
    )
    sapo_tau_neg: float = field(
        default=1.05,
        metadata={"help": "SAPO temperature for negative advantages"},
    )

    # Asynchronous RL
    recompute_logprob: bool = field(
        default=False,
        metadata={
            "help": "Recompute log probability and replace the log probability returned by inference."
        },
    )
    importance_sampling_level: str = field(
        default="token",
        metadata={
            "help": "Level at which to compute importance sampling ratios. 'token': per-token ratios (standard PPO). 'sequence': sequence-level geometric mean of per-token ratios (GSPO).",
            "choices": ["token", "sequence"],
        },
    )
    # Logging Agent Trajectories
    log_agent_stats: bool = field(
        default=False,
        metadata={"help": "Log statistics for agent trajectories"},
    )
    log_agent_stats_keys: list[str] = field(
        default_factory=lambda: [],
        metadata={"help": "Keys for logging agent trajectory statistics"},
    )
    # Others
    max_new_tokens: int = field(
        default=1024,
        metadata={"help": "Maximum number of new tokens to generate"},
    )


@dataclass
class PPOCriticConfig(TrainEngineConfig):
    """Configuration for PPO critic model, a subclass of a TrainEngine."""

    ppo_n_minibatches: int = field(
        default=4, metadata={"help": "Number of minibatches for each PPO update"}
    )
    eps_clip: float = field(
        default=0.5, metadata={"help": "Clipping factor for value loss"}
    )
    mask_no_eos_with_zero: bool = field(
        default=False,
        metadata={
            "help": "Mask truncated generations (no EOS token) and exclude from training"
        },
    )


def get_py_cmd(module: str, args: dict[str, Any]):
    # convert to flags
    cmd = ["python3", "-m", module]
    for k, v in args.items():
        if v is None or v is False or v == "" or (isinstance(v, list) and not v):
            continue
        flag = f"--{k.replace('_', '-')}"
        if v is True:
            cmd.append(flag)
        elif isinstance(v, list):
            cmd.append(flag)
            cmd.extend(map(str, v))
        else:
            cmd.append(flag)
            cmd.append(str(v))
    return cmd


@dataclass
class vLLMConfig:
    """Configuration for vLLM runtime. Refer to:
    https://docs.vllm.ai/en/stable/api/index.html for detailed documentation.
    """

    model: str = ""
    seed: int = 1
    skip_tokenizer_init: bool = False
    enforce_eager: bool = False
    dtype: str = "bfloat16"
    distributed_executor_backend: str = "mp"
    # original
    max_num_seqs: int = 256
    # kv_cache_type: str = "auto"
    block_size: int = 16
    swap_space: int = 4
    cpu_offload_gb: float = 0
    disable_sliding_window: bool = True
    # NOTE: Defaults max_model_len to 32k because a larger value
    # will enable chunked prefill in vLLM, which will cause
    # evalution performance degeneration.
    max_model_len: int | None = 32768
    enable_chunked_prefill: bool = False
    # NOTE: Setting enable_prefix_caching to False
    # because it will reuse the block after
    # model weights are updated. Using v0.7.2 reset_prefix_cache
    # will fix this issue.
    enable_prefix_caching: bool = False
    gpu_memory_utilization: float = 0.9
    worker_extension_cls: str = (
        "astraflow.train_worker.launcher.vllm.vllm_worker_extension.VLLMWorkerExtension"
    )
    enable_sleep_mode: bool = False
    uvicorn_log_level: str = "warning"
    enable_lora: bool = False
    lora_modules: str = ""

    @staticmethod
    def build_args(
        vllm_config: "vLLMConfig",
        tp_size: int,
        pp_size: int,
        host: str | None = None,
        port: int | None = None,
        dist_init_addr: str | None = None,
    ):
        args: dict = conf_as_dict(vllm_config)
        args = dict(
            # Model and tokenizer
            tokenizer=vllm_config.model,
            load_format="auto",
            trust_remote_code=True,
            tensor_parallel_size=tp_size,
            pipeline_parallel_size=pp_size,
            **args,
        )
        if port is not None:
            args["port"] = port
        if host is not None:
            args["host"] = host
        # handle lora modules separately
        lm = args.get("lora_modules")
        if lm:
            if isinstance(lm, str):
                lm = [lm]
            if isinstance(lm, (list, tuple)):
                try:
                    args["lora_modules"] = [
                        json.dumps(json.loads(s), separators=(",", ":")) for s in lm
                    ]
                except json.JSONDecodeError as e:
                    raise ValueError(f"Invalid JSON string in lora_modules: {e}") from e
        return args

    @staticmethod
    def build_cmd_from_args(args: dict[str, Any]):
        return get_py_cmd("astraflow.train_worker.launcher.vllm.vllm_server", args)

    @staticmethod
    def build_cmd(
        vllm_config: "vLLMConfig",
        tp_size: int,
        pp_size: int,
        host: str | None = None,
        port: int | None = None,
        dist_init_addr: str | None = None,
    ):
        args = vLLMConfig.build_args(
            vllm_config=vllm_config,
            tp_size=tp_size,
            pp_size=pp_size,
            host=host,
            port=port,
            dist_init_addr=dist_init_addr,
        )
        return vLLMConfig.build_cmd_from_args(args)


@dataclass
class SGLangConfig:
    """Configuration for SGLang runtime. Refer to:
    https://github.com/sgl-project/sglang for detailed documentation.
    """

    model_path: str = ""
    random_seed: int = 1
    skip_tokenizer_init: bool = False
    disable_cuda_graph: bool = False
    disable_radix_cache: bool = True
    disable_cuda_graph_padding: bool = False
    enable_nccl_nvls: bool = False
    disable_outlines_disk_cache: bool = False
    disable_custom_all_reduce: bool = False
    disable_overlap_schedule: bool = False
    enable_mixed_chunk: bool = False
    enable_dp_attention: bool = False
    enable_torch_compile: bool = False
    torch_compile_max_bs: int = 32
    cuda_graph_max_bs: int | None = None
    cuda_graph_bs: list[int] | None = None
    torchao_config: str = ""
    enable_nan_detection: bool = False
    enable_p2p_check: bool = False
    triton_attention_reduce_in_fp32: bool = False
    triton_attention_num_kv_splits: int = 8
    num_continuous_decode_steps: int = 1
    enable_memory_saver: bool = False
    allow_auto_truncate: bool = False
    attention_backend: str | None = "fa3"
    enable_multimodal: bool = False
    sampling_backend: str | None = None
    context_length: int | None = 32768
    mem_fraction_static: float | None = 0.9
    max_running_requests: int | None = None
    # NOTE: chunked_prefill_size is by default 8192 on GPUs with 80GB mem in SGLang,
    # but we disable it to avoid precision issues
    chunked_prefill_size: int | None = -1
    max_prefill_tokens: int = 32768
    schedule_policy: str = "lpm"
    schedule_conservativeness: float = 1.0
    cpu_offload_gb: int = 0
    dtype: str = "bfloat16"
    kv_cache_dtype: str = "auto"
    dp_size: int = 1  # only used for dp attention
    ep_size: int = 1
    # lora
    enable_lora: bool | None = None
    max_lora_rank: int | None = None
    lora_target_modules: list[str] | None = None
    lora_paths: list[str] | None = None
    max_loaded_loras: int = 1
    max_loras_per_batch: int = 1
    lora_backend: str = "triton"
    # logging
    log_level: str = "warning"
    log_level_http: str | None = "warning"
    log_requests: bool = False
    log_requests_level: int = 0
    show_time_cost: bool = False
    enable_metrics: bool = True  # Exports Prometheus-like metrics
    # The interval (in decoding iterations) to log throughput
    # and update prometheus metrics
    decode_log_interval: int = 1
    # Extra loader arguments
    # NOTE: These arguments will be parsed into a dict json-string
    # and passed as `model_loader_extra_config` to SGLang.
    enable_multithread_load: bool = False
    enable_fast_load: bool = False

    # Use staticmethod to make OmegaConf happy.
    @staticmethod
    def build_cmd(
        sglang_config: "SGLangConfig",
        tp_size,
        base_gpu_id,
        host: str | None = None,
        port: int | None = None,
        dist_init_addr: str | None = None,
        n_nodes: int = 1,
        node_rank: int = 0,
    ):
        args = SGLangConfig.build_args(
            sglang_config=sglang_config,
            tp_size=tp_size,
            base_gpu_id=base_gpu_id,
            host=host,
            port=port,
            dist_init_addr=dist_init_addr,
            n_nodes=n_nodes,
            node_rank=node_rank,
        )

        return SGLangConfig.build_cmd_from_args(args)

    @staticmethod
    def build_cmd_from_args(args: dict[str, Any]):
        return get_py_cmd("sglang.launch_server", args)

    @staticmethod
    def build_args(
        sglang_config: "SGLangConfig",
        tp_size: int,
        base_gpu_id: int,
        host: str | None = None,
        port: str | None = None,
        dist_init_addr: str | None = None,
        n_nodes: int = 1,
        node_rank: int = 0,
    ):
        # Map "all-linear" to "all"
        args: dict = conf_as_dict(sglang_config)
        if sglang_config.enable_multithread_load or sglang_config.enable_fast_load:
            if not pkg_version.is_version_equal("sglang", "0.5.2"):
                raise RuntimeError(
                    "Customized model loading requires exact SGLang version 0.5.2"
                )
            model_loader_extra_config = dict(
                enable_multithread_load=sglang_config.enable_multithread_load,
                enable_fast_load=sglang_config.enable_fast_load,
            )
            args["model_loader_extra_config"] = json.dumps(
                model_loader_extra_config, separators=(",", ":")
            )
        args.pop("enable_multithread_load", None)
        args.pop("enable_fast_load", None)
        # Map "all-linear" to "all"
        if "lora_target_modules" in args and args["lora_target_modules"]:
            args["lora_target_modules"] = [
                x.replace("-linear", "") for x in args["lora_target_modules"]
            ]
        from astraflow.train_worker.platforms import current_platform

        args = dict(
            # Model and tokenizer
            tokenizer_path=sglang_config.model_path,
            tokenizer_mode="auto",
            load_format="auto",
            trust_remote_code=True,
            device=current_platform.device_type,
            is_embedding=False,
            # Other runtime options
            tp_size=tp_size,
            # Because we have set CUDA_VISIBLE_DEVICES to a single GPU in each process
            base_gpu_id=base_gpu_id,
            nnodes=n_nodes,
            node_rank=node_rank,
            # initialization addresses and ports
            dist_init_addr=dist_init_addr,
            **args,
        )
        if host is not None:
            args["host"] = host
        if port is not None:
            args["port"] = port
        if not pkg_version.is_version_greater_or_equal("sglang", "0.4.9.post2"):
            raise RuntimeError("Needs sglang>=0.4.9.post2 to run the code.")
        if is_version_less("sglang", "0.4.10.post2"):
            args.pop("max_loaded_loras", None)
        return args


@dataclass
class InferenceEngineConfig:
    """Configuration for inference servers, including offpolicyness control."""

    experiment_name: str | None = None
    trial_name: str | None = None
    max_concurrent_rollouts: None | int = field(
        default=None,
        metadata={
            "help": "Maximum number of concurrent rollouts to "
            "the inference engine. Defaults to consumer_batch_size."
        },
    )
    queue_size: None | int = field(
        default=None,
        metadata={"help": "Input/Output queue size for async rollout."},
    )
    consumer_batch_size: int = field(
        default=16,
        metadata={"help": "Batch size for consuming rollouts from the queue."},
    )
    max_head_offpolicyness: int = field(
        default=0,
        metadata={
            "help": "Maximum off-policyness for the head. "
            "If the current version is more than this many versions behind, "
            "the request will not be accepted.",
        },
    )
    enable_rollout_tracing: bool = field(
        default=False,
        metadata={
            "help": "Whether to output verbose tracing messages for each generation request."
        },
    )
    check_trajectory_format: bool = field(
        default=False,
        metadata={
            "help": "Whether to check the format of produced trajectories of a customized workflow. Useful when debugging the workflow in isolation. Should be False during RL training."
        },
    )
    schedule_policy: str = field(
        default="round_robin",
        metadata={"help": "Request scheduling policy", "choices": ["round_robin"]},
    )
    setup_timeout: float = field(
        default=600.0,
        metadata={
            "help": "Timeout in seconds of connecting to remote servers or launching local servers."
        },
    )
    request_timeout: float = field(
        default=3600, metadata={"help": "Timeout for HTTP requests."}
    )
    request_retries: int = field(
        default=3, metadata={"help": "Number of retries for failed requests."}
    )
    pause_grace_period: float = field(
        default=0.0,
        metadata={
            "help": "The grace period after calling /pause_generation. Wait until all requests have been dropped."
        },
    )
    scheduling_strategy: SchedulingStrategy = field(
        default_factory=SchedulingStrategy,
        metadata={
            "help": "The scheduling strategy of this InferenceEngine, either separation or colocation."
        },
    )
    use_lora: bool = field(
        default=False,
        metadata={"help": "Whether to use LoRA. Should be same as actors LORA option."},
    )


@dataclass
class _Timer:
    freq_epochs: int | None = field(
        default=None,
        metadata={
            "help": "Trigger frequency in epochs. None disables epoch-based saving."
        },
    )
    freq_steps: int | None = field(
        default=None,
        metadata={
            "help": "Trigger frequency in steps. None disables step-based saving."
        },
    )
    freq_secs: int | None = field(
        default=None,
        metadata={
            "help": "Trigger frequency in seconds. None disables time-based saving."
        },
    )


@dataclass
class EvaluatorConfig(_Timer):
    """Configuration for evaluator scheduling."""

    eval_at_start: bool = field(
        default=True,
        metadata={
            "help": "If True, run a full evaluation pass at version=0 "
            "(before the first training step) on fresh runs. Ignored on "
            "recovery runs (start_step > 0)."
        },
    )


@dataclass
class SaverConfig(_Timer):
    """Configuration for model checkpoint saving scheduling and timing."""

    fileroot: str = MISSING


@dataclass
class RecoverConfig(_Timer):
    """Configuration for experiment recovery and fault tolerance."""

    fileroot: str = MISSING
    mode: str = field(
        default="disabled",
        metadata={
            "help": "Recovery mode for the launcher. "
            "Options: "
            "'disabled': Never recover from previous runs. "
            "'auto': Automatically recover from previous runs if recover info and checkpoints are available. "
            "'fault': Only recover from previous runs if the new run fails. "
            "'resume': Force to resume, raise an error if no recover info was found. Never resume if failed again."
        },
    )
    retries: int = field(
        default=3,
        metadata={"help": "Number of recovery retries (auto/fault modes only)."},
    )


@dataclass
class WandBConfig:
    """Configuration for Weights & Biases experiment tracking."""

    mode: str = "disabled"
    wandb_dir: str | None = None
    wandb_base_url: str = ""
    wandb_api_key: str = ""
    entity: str | None = None
    project: str | None = None
    name: str | None = None
    job_type: str | None = None
    group: str | None = None
    notes: str | None = None
    tags: list[str] | None = None
    config: dict | None = None
    id_suffix: str | None = "train"


@dataclass
class SwanlabConfig:
    """Configuration for SwanLab experiment tracking and monitoring."""

    project: str | None = None
    name: str | None = None
    config: dict | None = None
    logdir: str | None = None
    mode: str | None = "disabled"
    api_key: str | None = os.getenv("SWANLAB_API_KEY", None)


@dataclass
class TensorBoardConfig:
    """Configuration for TensorBoard logging and visualization."""

    path: str | None = None


@dataclass
class StatsLoggerConfig:
    """Configuration for experiment statistics logging and tracking services."""

    experiment_name: str = MISSING
    trial_name: str = MISSING
    fileroot: str = MISSING
    wandb: WandBConfig = field(
        default_factory=WandBConfig,
        metadata={"help": "Weights & Biases configuration."},
    )
    swanlab: SwanlabConfig = field(
        default_factory=SwanlabConfig,
        metadata={"help": "SwanLab configuration."},
    )
    tensorboard: TensorBoardConfig = field(
        default_factory=TensorBoardConfig,
        metadata={"help": "TensorBoard configuration. Only 'path' field required."},
    )


@dataclass
class SessionTracerConfig:
    """Configuration for per-session lifecycle tracing."""

    enabled: bool = field(
        default=False,
        metadata={
            "help": (
                "Enable per-session lifecycle tracing alongside perf events. "
                "When true, session metadata is captured to sessions.jsonl."
            )
        },
    )
    flush_threshold: int = field(
        default=256,
        metadata={
            "help": (
                "Flush session trace records once this many entries are ready. "
                "Values <= 0 fall back to 1."
            )
        },
    )


@dataclass
class PerfTracerConfig:
    """Configuration for perf tracer emission."""

    experiment_name: str = MISSING
    trial_name: str = MISSING
    fileroot: str = MISSING
    enabled: bool = field(
        default=False,
        metadata={
            "help": (
                "Explicitly enable or disable perf tracing. Set to true to capture perf traces."
            )
        },
    )
    save_interval: int = field(
        default=1,
        metadata={
            "help": (
                "Flush trace events to disk every N calls to save(step=...). "
                "A value of 1 writes on every step; values <= 0 fall back to 1."
            )
        },
    )
    profile_steps: list[int] | None = field(
        default=None,
        metadata={
            "help": (
                "List of step numbers at which to capture detailed profiling traces. "
                "If None, no detailed profiling traces are captured."
            )
        },
    )
    session_tracer: SessionTracerConfig | None = field(
        default=None,
        metadata={"help": "Session tracing configuration."},
    )


@dataclass
class NameResolveConfig:
    """Configuration for distributed name resolution and service discovery."""

    type: str = field(
        default="nfs",
        metadata={
            "help": "Type of the distributed KV store for name resolving.",
            "choices": ["nfs", "etcd3", "ray"],
        },
    )
    nfs_record_root: str = field(
        default="/tmp/astraflow/name_resolve",
        metadata={
            "help": "Record root for NFS name resolving. Should be available on all nodes."
        },
    )
    etcd3_addr: str = field(
        default="localhost:2379", metadata={"help": "Address of the ETCD3 server."}
    )
    ray_actor_name: str = field(
        default="ray_kv_store",
        metadata={"help": "Name of the distributed Ray KV store."},
    )


@dataclass
class ClusterSpecConfig:
    """Configuration for cluster specification and distributed computing setup."""

    name_resolve: NameResolveConfig = field(
        default_factory=NameResolveConfig,
        metadata={"help": "Name resolving configuration."},
    )
    cluster_name: str = field(
        default="local",
        metadata={"help": "Name of the cluster. Used to set specific environs."},
    )
    fileroot: str = field(
        default="/tmp/astraflow/",
        metadata={
            "help": "Root for logs and checkpoints. Should be available on all nodes."
        },
    )
    n_nodes: int = field(
        default=32,
        metadata={
            "help": "The size of the cluster. Used to decide slurm hostname suffix."
        },
    )
    n_gpus_per_node: int = field(
        default=8,
        metadata={"help": "Number of GPUs per node (physical)."},
    )


@dataclass
class SchedulerConfig:
    """Configuration for worker scheduling. Used in the single-controller mode. Experimental."""

    endpoint: str = field(default="http://localhost:8081")
    deploy_mode: str = field(default="separation")
    functioncall_service_domain: str = field(default="http://localhost:8080")
    reward_functioncall_config: dict = field(default_factory=dict)
    reward_model_path: str = field(default="")
    reward_model_service_url: str = field(default="http://localhost:30000/classify")


@dataclass
class SlurmLauncherConfig:
    """Configuration for launching the training jobs with Slurm."""

    srun_additional_args: str = field(
        default="--mpi=pmi2 -K --chdir $PWD",
        metadata={"help": "Additional arguments to pass to the srun command."},
    )
    additional_bash_cmds: list[str] | None = field(
        default=None,
        metadata={
            "help": "Additional bash commands to setup the container before running "
            "the torchrun command."
        },
    )
    container_type: str = field(
        default="apptainer",
        metadata={
            "help": "Type of containers used in slurm",
            "choices": ["apptainer", "none"],
        },
    )
    mount: str = field(
        default="/storage:/storage", metadata={"help": "Mount path for slurm."}
    )
    trainer_image: str | None = field(
        default=None, metadata={"help": "slurm image for trainers."}
    )
    inference_server_image: str | None = field(
        default=None, metadata={"help": "slurm image for LLM inference."}
    )


@dataclass
class LauncherConfig:
    """Configuration for launching the LLM server and trainer processes."""

    inference_server_cpus_per_gpu: int = field(
        default=4,
        metadata={"help": "Number of CPUs allocated per GPU for inference server."},
    )
    inference_server_mem_per_gpu: int = field(
        default=32 * 1024,
        metadata={"help": "Memory allocated per GPU for inference server in MB."},
    )
    trainer_cpus_per_gpu: int = field(
        default=4,
        metadata={"help": "Number of CPUs allocated per GPU for training."},
    )
    trainer_mem_per_gpu: int = field(
        default=32 * 1024,
        metadata={"help": "Memory allocated per GPU for training in MB."},
    )
    inference_server_env_vars: str = field(
        default="",
        metadata={
            "help": "Environment variables for inference server, separated by commas. "
            "Example: 'ENV1=val1,ENV2=val2'."
        },
    )
    trainer_env_vars: str = field(
        default="",
        metadata={
            "help": "Environment variables for training, separated by commas. "
            "Example: 'ENV1=val1,ENV2=val2'."
        },
    )
    slurm: SlurmLauncherConfig = field(
        default_factory=SlurmLauncherConfig,
        metadata={"help": "Slurm launcher configuration."},
    )


@dataclass
class BufferConfig:
    """Configuration for rollout buffer."""

    size: int = field(
        default=65536,
        metadata={"help": "Maximum number of examples to store in the buffer."},
    )
    debug: bool = field(
        default=False,
        metadata={"help": "If True, enable debug printing in the rollout buffer."},
    )
    filter_function: str | None = field(
        default=None,
        metadata={
            "help": "Name of the filter function from the registry. "
            "Available filters: 'keep_all', 'filter_zero_adv'. "
            "If None, uses KeepAllFilter (accepts all examples)."
        },
    )
    max_staleness: int | None = field(
        default=None,
        metadata={
            "help": "Maximum allowed staleness (version gap) for filtering examples. "
            "If current_version - min_version > max_staleness, the example will be skipped. "
            "If None, staleness filtering is disabled. "
            "If not set in buffer config, falls back to rollout.max_head_offpolicyness."
        },
    )
    replay_size: int | None = field(
        default=None,
        metadata={
            "help": "Maximum number of examples to store in the replay buffer. "
            "If None, defaults to the main buffer size."
        },
    )
    replay_ratio: float = field(
        default=0.0,
        metadata={
            "help": "Fraction of each training batch to draw from the replay buffer. "
            "Must be between 0 and 1."
        },
    )
    selection_fn: str | None = field(
        default="latest",
        metadata={
            "help": "Replay selection strategy name. Available selections: 'latest'. "
            "Defaults to 'latest'."
        },
    )
    weight_pull_interval: int = field(
        default=2,
        metadata={
            "help": "Pull new weights from WeightController every N training steps. "
            "Only used in decoupled mode (WeightController). "
            "A new weight version is pushed to RaaS when "
            "latest_version - last_pulled_version >= weight_pull_interval."
        },
    )


@dataclass
class BaseExperimentConfig:
    """Base configuration class for all experiment types with common settings."""

    # NOTE: we need this unified config class because different experiments
    # have different config structures, e.g., GRPO has two engine configs,
    # but SFT only has a single one. We use subclasses to represent these structures.
    experiment_name: str = field(
        default=MISSING,
        metadata={"help": "Name of the experiment (no '_' or '/'). Required."},
    )
    trial_name: str = field(
        default=MISSING,
        metadata={"help": "Name of the trial (no '-' or '/'). Required."},
    )
    cluster: ClusterSpecConfig = field(
        default_factory=ClusterSpecConfig,
        metadata={"help": "Cluster specification. Mainly used by slurm."},
    )
    allocation_mode: Any = field(
        default="",
        metadata={"help": "Engine allocation config. Dict (engine section) or string (legacy). "},
    )
    seed: int = field(default=1, metadata={"help": "Random seed for reproducibility."})
    enable_offload: bool = field(
        default=False,
        metadata={
            "help": "Whether to enable training offload using torch_memory_saver. "
            "This requires setting up the environment for TMS (e.g., via LD_PRELOAD)."
        },
    )
    total_train_epochs: int = field(
        default=1, metadata={"help": "Total number of epochs to train the model."}
    )
    total_train_steps: int | None = field(
        default=None,
        metadata={
            "help": "Terminate training after this number of steps. "
            "For benchmarking purposes only. None indicates normal training."
        },
    )
    total_train_n_seqs: int | None = field(
        default=None,
        metadata={
            "help": "Terminate training after consuming this number of samples. "
            "For benchmarking purposes only. None indicates normal training."
        },
    )
    tokenizer_path: str = field(
        default="",
        metadata={"help": "Path to the tokenizer."},
    )

    saver: SaverConfig = field(default_factory=SaverConfig)
    evaluator: EvaluatorConfig = field(default_factory=EvaluatorConfig)
    stats_logger: StatsLoggerConfig = field(default_factory=StatsLoggerConfig)
    perf_tracer: PerfTracerConfig | None = field(
        default=None,
        metadata={"help": "Performance tracer configuration. None means disabled."},
    )
    recover: RecoverConfig = field(default_factory=RecoverConfig)

    sglang: SGLangConfig = field(default_factory=SGLangConfig)
    vllm: vLLMConfig = field(default_factory=vLLMConfig)
    launcher: LauncherConfig = field(default_factory=LauncherConfig)

    scheduler: SchedulerConfig = field(default_factory=SchedulerConfig)

    buffer: BufferConfig = field(default_factory=BufferConfig)


@dataclass
class PPOConfig(BaseExperimentConfig):
    """Configuration for Proximal Policy Optimization (PPO) reinforcement learning experiments."""

    gconfig: GenerationHyperparameters = field(
        default_factory=GenerationHyperparameters
    )
    rollout: InferenceEngineConfig = field(default_factory=InferenceEngineConfig)
    actor: PPOActorConfig = field(default_factory=PPOActorConfig)
    ref: PPOActorConfig | None = field(default=None)
    critic: PPOCriticConfig | None = field(default=None)
    weight_transfer_mode: str = field(
        default="tcp",
        metadata={
            "help": (
                "Weight transfer mode after each PPO step. "
                "'tcp' (default): non-blocking TCP transfer via sender_agent — "
                "RaaS pulls weights asynchronously; training continues immediately. "
                "NCCL mode has been removed."
            )
        },
    )
    weight_transfer_strategies: str = field(
        default="full",
        metadata={
            "help": (
                "Weight transfer strategy. "
                "'full': send complete weights each step. "
                "'delta': send sparse diffs (full is always available as fallback)."
            )
        },
    )
    model_id: str | None = field(
        default=None,
        metadata={
            "help": (
                "Model identifier for multi-model training (e.g. 'model0', 'model1'). "
                "When set, the trainer participates in multi-model version barriers "
                "and RaaS weight updates are scoped to this model. "
                "When None (default), single-model mode is used."
            )
        },
    )
    agent_name: str = field(
        default="default",
        metadata={
            "help": (
                "Agent name for connecting to the AstraFlow HTTP service. "
                "Must match a key in the AstraFlow service's `agents` dict. "
                "In multi-model training, each trainer uses a distinct agent_name "
                "(e.g. 'model0', 'model1') to get its own batch and version tracking."
            )
        },
    )

    # Shared generation parameters — set once in `shared` YAML section.
    # Propagated to actor.group_size, actor.reward_norm.group_size, and
    # actor.temperature in __post_init__ if those are not explicitly set.
    n_samples: int = field(
        default=1,
        metadata={"help": "Number of samples per prompt. Propagates to actor.group_size and actor.reward_norm.group_size."},
    )
    temperature: float = field(
        default=1.0,
        metadata={"help": "Sampling temperature. Propagates to actor.temperature."},
    )
    train_batch_size: int = field(
        default=256,
        metadata={
            "help": "Number of examples (sequences) per training step. "
            "Global batch size across all data-parallel ranks. "
            "In AstraFlow v2, this is sent to the HTTP service to control batch size."
        },
    )

    def __post_init__(self):
        # Propagate shared generation params to actor config.
        # Only fill in when the actor field is still at its default (None),
        # meaning the user did not explicitly set it in the YAML.
        if self.actor is not None:
            if self.actor.group_size is None:
                self.actor.group_size = self.n_samples
            if self.actor.temperature is None:
                self.actor.temperature = self.temperature
            if self.actor.reward_norm is not None:
                if self.actor.reward_norm.group_size is None:
                    self.actor.reward_norm.group_size = self.n_samples
        # Ref model also needs temperature for logprob computation.
        if self.ref is not None:
            if self.ref.temperature is None:
                self.ref.temperature = self.temperature


@dataclass
class GRPOConfig(PPOConfig):
    """A dummy place holder of GRPO config for backward compatibility."""

    offline_data_path: str = field(
        default="data-data",
        metadata={"help": "Base directory containing offline datasets."},
    )


def parse_cli_args(argv: list[str]):
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--config", help="Path to the experiment configuration file", required=True
    )
    parser.add_argument(
        "--trainer",
        type=str,
        default=None,
        help="Trainer section key (e.g. 'trainer_model0'). Default: 'trainer'.",
    )
    # The first argument might be the path to a training script,
    # which should be ignored by the argument parser.
    if argv and argv[0].endswith(".py"):
        argv = argv[1:]
    args, overrides = parser.parse_known_args(argv)
    config_file = Path(args.config).absolute()
    assert config_file.exists(), f"Config file {config_file} does not exist."

    from astraflow.core.config.loader import load_and_merge_configs, load_trainer_config

    raw = load_and_merge_configs([str(config_file)])
    trainer_key = args.trainer or "trainer"
    trainer_dict = load_trainer_config(raw, trainer_key=trainer_key)

    # Apply CLI dotted-key overrides
    for override in overrides:
        if "=" in override:
            key, val = override.split("=", 1)
            parts = key.split(".")
            d = trainer_dict
            for p in parts[:-1]:
                if p not in d:
                    d[p] = {}
                d = d[p]
            try:
                parsed = yaml.safe_load(val)
            except Exception:
                parsed = val
            d[parts[-1]] = parsed

    cfg = OmegaConf.create(trainer_dict)
    return cfg, config_file


def to_structured_cfg(cfg, config_cls):
    # Merge with the default configuration.
    # The yaml and commandline can omit some default values defined in python dataclasses.
    default_cfg = OmegaConf.structured(config_cls)
    # Strip top-level keys from the YAML that don't exist on the dataclass
    # (e.g. shared scalars like model_path, dtype that are only used for
    # ${...} substitution and shouldn't be stored in the config object).
    known_keys = set(default_cfg.keys())
    unknown = [k for k in cfg if k not in known_keys]
    if unknown:
        OmegaConf.set_struct(cfg, False)
        for k in unknown:
            del cfg[k]
    cfg = OmegaConf.merge(default_cfg, cfg)
    return cfg


def load_expr_config(argv: list[str], config_cls: type[ConfigT]) -> tuple[ConfigT, str]:
    cfg, config_file = parse_cli_args(argv)
    cfg = to_structured_cfg(cfg, config_cls=config_cls)
    cfg = OmegaConf.to_object(cfg)
    assert isinstance(cfg, config_cls)
    # Setup environment

    name_resolve.reconfigure(cfg.cluster.name_resolve)

    from astraflow.train_worker.utils.stats_logger import StatsLogger

    # Save configuration as yaml
    if os.getenv("RANK", "0") == "0":
        save_config(cfg, StatsLogger.get_log_path(cfg.stats_logger))

    return cfg, str(config_file)


def conf_as_dict(cfg):
    if isinstance(cfg, (OmegaConf, DictConfig)):
        return OmegaConf.to_container(cfg, resolve=True)
    return asdict(cfg)


def save_config(cfg, log_dir):
    os.makedirs(log_dir, exist_ok=True)
    config_save_path = os.path.join(log_dir, "config.yaml")
    with open(config_save_path, "w") as f:
        config_dict: dict = asdict(cfg)
        yaml.dump(
            config_dict,
            f,
            default_flow_style=False,
            sort_keys=False,
        )