# Copyright (c) ModelScope Contributors. All rights reserved.
import json
import megatron.core
import os
import torch
from dataclasses import dataclass, field, fields
from mcore_bridge.model import get_mcore_model_type, get_model_meta
from megatron.core import mpu
from megatron.core.transformer.enums import AttnBackend
from packaging import version
from transformers.utils import is_torch_npu_available
from transformers.utils.versions import require_version
from typing import Any, Dict, List, Literal, Optional, Tuple, Union

from swift.arguments import ModelArguments
from swift.megatron.utils import initialize_megatron
from swift.model import get_model_info_meta
from swift.utils import get_dist_setting, get_logger, json_parse_to_dict

mcore_016 = version.parse(megatron.core.__version__) >= version.parse('0.16.0rc0')
logger = get_logger()


@dataclass
class RLHFMegatronArgumentsMixin:
    rlhf_type: Literal['dpo', 'kto', 'grpo', 'gkd', 'rm'] = None
    loss_type: Optional[str] = None  # rlhf / plugins
    mcore_ref_model: Optional[str] = None
    mcore_ref_adapter: Optional[str] = None

    beta: Optional[float] = None
    rpo_alpha: Optional[float] = None
    reference_free: bool = False
    label_smoothing: float = 0.
    f_divergence_type: str = 'reverse_kl'

    # kto
    desirable_weight: float = 1.
    undesirable_weight: float = 1.
    calculate_KL: Optional[bool] = None

    # rm
    center_rewards_coefficient: Optional[float] = None

    # gkd
    teacher_model: Optional[str] = field(default=None)
    teacher_model_type: Optional[str] = field(default=None)
    teacher_model_revision: Optional[str] = field(default=None)
    _teacher_use_disable_adapter: bool = False
    teacher_model_server: Optional[str] = field(
        default=None,
        metadata={
            'help':
            'URL of the teacher model server (e.g., http://localhost:8000). '
            'When set, teacher logprobs are fetched via API instead of loading a local model. '
            'Supports multi-teacher via JSON.'
        })
    teacher_tag_key: str = field(
        default='dataset', metadata={'help': 'Column name for multi-teacher routing. Default "dataset".'})
    gkd_logits_topk: Optional[int] = None
    lmbda: float = 0.5  # On-policy probability: with prob lmbda, use student-generated responses
    seq_kd: bool = False  # Deprecated
    offload_teacher_model: bool = False  # Offload teacher model to CPU to save GPU memory
    sft_alpha: float = 0.0  # Weight for SFT loss in GKD (0 = pure JSD, >0 = JSD + sft_alpha * SFT)

    # OPD-RL (On-Policy Distillation as RL): a teacher (teacher_model / teacher_model_server)
    # on a GRPO run turns it into OPD-RL, injecting teacher KL as the advantage.
    teacher_kl_coef: float = 1.0

    # grpo/gkd
    temperature: float = 0.9  # Temperature for sampling and loss computation

    # grpo
    generation_batch_size: Optional[int] = None
    steps_per_generation: Optional[int] = None
    num_generations: int = 8
    num_generations_eval: Optional[int] = None
    max_completion_length: int = 512
    # GSPO https://arxiv.org/abs/2507.18071
    importance_sampling_level: Literal['token', 'sequence', 'sequence_token'] = 'token'

    # SAPO https://arxiv.org/abs/2511.20347
    # Temperature parameters for soft adaptive gate
    tau_pos: float = 1.0
    tau_neg: float = 1.05

    # REAL https://arxiv.org/abs/2602.05630
    real_tau: float = 0.5

    # FIPO https://arxiv.org/abs/2603.19835
    fipo_decay_rate: float = 32.0
    fipo_clip_range: Optional[float] = 0.2
    fipo_clip_high_only: bool = True
    fipo_safety_threshold: Optional[float] = 4.0

    epsilon: float = 0.2
    epsilon_high: Optional[float] = None
    delta: Optional[float] = None
    top_k: int = -1
    top_p: float = 1.0
    repetition_penalty: float = 1.

    use_vllm: bool = True
    use_ray: bool = False
    vllm_mode: Optional[Literal['server', 'colocate']] = None

    vllm_enable_prefix_caching: bool = True
    vllm_gpu_memory_utilization: float = 0.9
    vllm_tensor_parallel_size: int = 1
    vllm_max_model_len: Optional[int] = None
    vllm_enforce_eager: bool = False
    vllm_limit_mm_per_prompt: Optional[Union[dict, str]] = None  # '{"image": 5, "video": 2}'
    vllm_disable_cascade_attn: bool = False
    vllm_max_num_seqs: Optional[int] = None
    vllm_mm_processor_cache_gb: Optional[float] = None
    vllm_engine_kwargs: Optional[Union[dict, str]] = None
    vllm_enable_lora: bool = False

    sleep_level: Literal[0, 1, 2] = 0
    offload_optimizer: bool = False
    offload_model: bool = False
    offload_bridge: bool = False

    vllm_server_base_url: Optional[List[str]] = None
    vllm_server_host: Optional[List[str]] = None
    vllm_server_port: List[int] = field(default_factory=lambda: [8000])
    vllm_server_timeout: float = 240.0
    vllm_server_group_port: Optional[List[int]] = None

    reward_funcs: List[str] = field(default_factory=list)
    reward_weights: List[float] = None
    # see details in swift/rewards/orm.py
    # cosine reward, https://arxiv.org/abs/2502.03373
    cosine_min_len_value_wrong: float = -0.5  # r^w_0 in paper, Reward for wrong answers with zero completion length.
    cosine_max_len_value_wrong: float = 0.0  # r^w_L in paper, Reward for wrong answers with max completion length.
    cosine_min_len_value_correct: float = 1.0  # r^c_0 in paper, Reward for correct answers with zero completion length.
    cosine_max_len_value_correct: float = 0.5  # r^c_L in paper, Reward for correct answers with max completion length.
    cosine_max_len: Optional[int] = None  # Lmax in paper, default equal to max_completion_length
    # repetition penalty, https://arxiv.org/abs/2502.03373
    repetition_n_grams: int = 3
    repetition_max_penalty: float = -1.0
    # soft_overlong, https://arxiv.org/abs/2503.14476
    soft_max_length: Optional[int] = None
    soft_cache_length: Optional[int] = None
    # DAPO, https://arxiv.org/abs/2503.14476
    dynamic_sample: bool = False
    max_resample_times: int = 3
    overlong_filter: bool = False

    # Dr. GRPO, https://arxiv.org/abs/2503.20783
    # GDPO: normalize each reward function separately
    scale_rewards: Literal['none', 'group', 'batch', 'gdpo'] = 'group'

    # RLOO / REINFORCE++
    advantage_estimator: Literal['grpo', 'rloo', 'reinforce_plus_plus'] = 'grpo'
    kl_in_reward: bool = False

    wandb_log_unique_prompts: Optional[bool] = None
    log_completions: bool = False

    rollout_importance_sampling_mode: Optional[Literal['token_truncate', 'token_mask', 'sequence_truncate',
                                                       'sequence_mask']] = None
    rollout_importance_sampling_threshold: float = 2.0
    log_rollout_offpolicy_metrics: bool = False
    # Off-Policy Sequence Masking: mask out sequences that deviate too much from rollout policy
    # If set, compute mean(rollout_per_token_logps - per_token_logps) per sequence,
    # and mask sequences where this delta > threshold AND advantage < 0
    # Falls back to old_per_token_logps if rollout_per_token_logps is not available
    off_policy_sequence_mask_delta: Optional[float] = None

    # entropy
    log_entropy: bool = False
    # Beyond the 80/20 Rule, https://arxiv.org/abs/2506.01939
    top_entropy_quantile: float = 1.0

    router_replay_mode: Literal['disabled', 'R2', 'R3'] = 'disabled'

    # ───────────────────────────  Not Supported Yet  ───────────────────────────

    # reward model
    reward_model: Optional[List[str]] = None
    reward_model_plugin: Optional[List[str]] = None
    # sync ref model
    sync_ref_model: bool = False
    ref_model_sync_steps: int = 512
    ref_model_mixup_alpha: float = 0.6

    async_generate: bool = False

    move_model_batches: Optional[int] = None

    # multi turn
    multi_turn_scheduler: Optional[str] = None
    max_turns: Optional[int] = None
    completion_length_limit_scope: Literal['total', 'per_round'] = 'per_round'
    vllm_server_pass_dataset: bool = False
    use_gym_env: Optional[bool] = None
    gym_env: Optional[str] = None

    num_iterations: int = 1

    # dataset
    dataset_shuffle: Optional[bool] = True

    def _init_kto(self):
        if self.calculate_KL is None:
            # Not all losses require a KL calculation
            self.calculate_KL = True
            if self.loss_type in ['apo_zero_unpaired']:
                self.calculate_KL = False

    def __post_init__(self):
        if self.rlhf_type is None:
            return
        default_loss_type = {'kto': 'kto', 'dpo': 'sigmoid', 'grpo': 'grpo'}
        default_beta = {'gkd': 0.5, 'grpo': 0.04}
        if self.beta is None:
            self.beta = default_beta.get(self.rlhf_type, 0.1)
        if self.loss_type is None:
            self.loss_type = default_loss_type.get(self.rlhf_type)
        if self.rlhf_type == 'kto':
            self._init_kto()
        if self.rlhf_type == 'grpo':
            self._init_grpo()
            if self.cosine_max_len is None:
                self.cosine_max_len = self.max_completion_length
            if self.vllm_limit_mm_per_prompt is not None:
                self.vllm_limit_mm_per_prompt = json_parse_to_dict(self.vllm_limit_mm_per_prompt)
        # Teacher setup is identical for GKD and GRPO (OPD-RL): a teacher_model / server /
        # same-model LoRA self-distillation all flow through the same detection. GKD also
        # allows dynamic self-distillation (no teacher at all); GRPO without a teacher is
        # plain RL, so only resolve a teacher for GRPO when one is configured.
        if self.rlhf_type == 'gkd' or (self.rlhf_type == 'grpo' and
                                       (self.teacher_model is not None or self.teacher_model_server is not None)):
            self._check_teacher()
            if self.rlhf_type == 'grpo':
                self._check_opd_rl()
        if self.rlhf_type == 'gkd':
            # GKD-specific: the API path only returns top-k logprobs, so gkd_logits_topk is required.
            if self.teacher_model_server is not None and self.gkd_logits_topk is None:
                raise ValueError('gkd_logits_topk is required when using teacher_model_server')

            # Validate gkd_logits_topk
            if self.gkd_logits_topk is not None and self.gkd_logits_topk <= 0:
                raise ValueError(f'gkd_logits_topk must be a positive integer, got {self.gkd_logits_topk}')

            # seq_kd (teacher-generated responses) is not implemented; raise early.
            if self.seq_kd:
                raise NotImplementedError('seq_kd=True (Sequential KD with teacher generation) is not supported.')

            self.num_generations = 1
            self._init_generation_batch_params()

        if self.rlhf_type in ['grpo', 'gkd']:
            self.vllm_engine_kwargs = json_parse_to_dict(self.vllm_engine_kwargs)
            if self.use_vllm and os.getenv('SWIFT_AUDIO_LOAD_BACKEND') is None:
                # align with vLLM audio load backend
                os.environ['SWIFT_AUDIO_LOAD_BACKEND'] = 'soundfile_pyav'

    @staticmethod
    def resolve_generation_batch_size(
        generation_batch_size: Optional[int],
        steps_per_generation: Optional[int],
        global_batch_size: int,
        num_generations: int,
    ) -> Tuple[int, int]:
        """Resolve generation_batch_size and steps_per_generation from user inputs.

        Returns (generation_batch_size, steps_per_generation).
        """
        if generation_batch_size is None and steps_per_generation is None:
            steps_per_generation = 1
            generation_batch_size = global_batch_size * steps_per_generation
        elif generation_batch_size is not None and steps_per_generation is not None:
            expected = global_batch_size * steps_per_generation
            if generation_batch_size != expected:
                raise ValueError(f'generation_batch_size ({generation_batch_size}) must equal '
                                 f'global_batch_size ({global_batch_size}) * steps_per_generation '
                                 f'({steps_per_generation}) = {expected}.')
        elif generation_batch_size is not None:
            if generation_batch_size % global_batch_size != 0:
                raise ValueError(f'generation_batch_size ({generation_batch_size}) '
                                 f'must be divisible by global_batch_size ({global_batch_size})')
            steps_per_generation = generation_batch_size // global_batch_size
        else:
            generation_batch_size = global_batch_size * steps_per_generation

        if steps_per_generation <= 0:
            raise ValueError(f'steps_per_generation must be > 0, got {steps_per_generation}.')

        if num_generations > 1 and generation_batch_size % num_generations != 0:
            raise ValueError(f'generation_batch_size ({generation_batch_size}) must be divisible by '
                             f'num_generations ({num_generations}).')

        return generation_batch_size, steps_per_generation

    @staticmethod
    def validate_batch_dp_alignment(
        generation_batch_size: int,
        num_generations: int,
        dp_size: int,
        micro_batch_size: int,
        world_size: int,
    ) -> None:
        """Validate that batch parameters are correctly aligned with DP parallelism.

        Reusable across both worker-side (torch.distributed) and pipeline-side
        (Ray config-based dp estimation) contexts.
        """
        num_rollout_prompt = generation_batch_size // num_generations
        if num_rollout_prompt % dp_size != 0:
            raise ValueError(f'num_rollout_prompt ({num_rollout_prompt}) = generation_batch_size '
                             f'({generation_batch_size}) // num_generations ({num_generations}) '
                             f'must be divisible by dp_size ({dp_size}).')

        per_device_num_rollout_prompt = num_rollout_prompt // dp_size
        if per_device_num_rollout_prompt < 1:
            raise ValueError(f'per_device_num_rollout_prompt ({per_device_num_rollout_prompt}) must be >= 1, '
                             f'please adjust generation_batch_size/steps_per_generation/num_generations.')

        if per_device_num_rollout_prompt % micro_batch_size != 0:
            raise ValueError(f'Per-device rollout prompt count ({per_device_num_rollout_prompt}) = '
                             f'(generation_batch_size ({generation_batch_size}) // '
                             f'num_generations ({num_generations})) // dp_size ({dp_size}) '
                             f'must be divisible by micro_batch_size ({micro_batch_size}).')

        per_device_generation_batch_size = generation_batch_size // world_size
        if per_device_generation_batch_size < 1:
            raise ValueError(f'per_device_generation_batch_size ({per_device_generation_batch_size}) must be >= 1.')

    def _init_generation_batch_params(self):
        self.generation_batch_size, self.steps_per_generation = self.resolve_generation_batch_size(
            self.generation_batch_size, self.steps_per_generation, self.global_batch_size, self.num_generations)

        if torch.distributed.is_initialized():
            world_size = torch.distributed.get_world_size()
            dp_size = world_size // (
                self.pipeline_model_parallel_size * self.tensor_model_parallel_size * self.context_parallel_size)
            self.validate_batch_dp_alignment(self.generation_batch_size, self.num_generations, dp_size,
                                             self.micro_batch_size, world_size)
            self.per_device_generation_batch_size = self.generation_batch_size // world_size

    def _check_teacher(self):
        """Resolve the teacher (shared by GKD and GRPO/OPD-RL).

        Detects the three teacher modes and sets ``_teacher_use_disable_adapter``:
          - separate teacher_model / teacher_model_server,
          - same-model LoRA self-distillation (disable_adapter, no extra model),
          - dynamic self-distillation (no teacher -> teacher == current student weights).
        """
        if self.teacher_model is not None and self.teacher_model_server is not None:
            raise ValueError('setting both `teacher_model` and `teacher_model_server` is not supported.')

        # Fail fast: the Ray pipeline only supports a colocated teacher_model (see
        # swift/ray/megatron/grpo_trainer.py), so reject teacher_model_server at parse time.
        if self.use_ray and self.teacher_model_server is not None:
            raise ValueError('teacher_model_server is not supported with use_ray')

        # Validate teacher_model_server: accept single URL or JSON multi-teacher config.
        if self.teacher_model_server is not None:
            from swift.rlhf_trainers.gkd_helpers import parse_teacher_model_server
            parse_teacher_model_server(self.teacher_model_server)

        self._teacher_use_disable_adapter = False
        if self.teacher_model is not None and self.teacher_model == self.model:
            if self.tuner_type == 'lora':
                logger.info('LoRA + same teacher_model: using disable_adapter() for fixed teacher (no extra model).')
                self._teacher_use_disable_adapter = True
                self.teacher_model = None
                self.teacher_model_dir = None
            # Full training + same teacher_model: a separate frozen copy is loaded as the fixed teacher.

    def _check_opd_rl(self):
        """Fail-fast OPD-RL (teacher distillation on Megatron GRPO) parameter compatibility.

        Mirrors ``RLHFArguments._check_opd_rl``: the teacher signal is a *per-token* advantage, so
        features that need a *per-sequence* advantage (sign-based) or reward variance are rejected.
        Called after ``_init_grpo`` so ``loss_type`` / ``scale_rewards`` are already resolved.
        """
        if self.loss_type in ['real', 'fipo']:
            raise ValueError(f'OPD-RL (teacher) does not support loss_type={self.loss_type!r} '
                             '(it needs a per-sequence advantage). Use grpo/bnpo/dr_grpo/dapo/cispo/sapo.')
        if self.off_policy_sequence_mask_delta is not None:
            raise ValueError('OPD-RL (teacher) does not support off_policy_sequence_mask_delta '
                             '(it needs a per-sequence advantage).')
        if not self.reward_funcs:
            if self.dynamic_sample:
                raise ValueError('dynamic_sample requires reward_funcs (it filters groups by reward std); '
                                 'pure OPD-RL distillation has no reward variance.')
            if self.scale_rewards == 'gdpo':
                raise ValueError("scale_rewards='gdpo' requires reward_funcs; pure OPD-RL distillation has none.")

    def _init_grpo(self):

        def _check_not_supported():
            if self.async_generate:
                raise ValueError('async_generate is not supported for Megatron GRPO right now')
            if self.sync_ref_model:
                raise ValueError('sync_ref_model is not supported for Megatron GRPO right now')
            if self.num_iterations > 1:
                raise ValueError('num_iterations > 1 is not supported for Megatron GRPO right now')

            if self.loss_type == 'real':
                assert self.micro_batch_size % self.num_generations == 0, \
                    (f'"REAL loss requires that the training micro_batch_size ({self.micro_batch_size}) '
                     f'is a multiple of num_generations ({self.num_generations}). Please adjust your batch parameters.')

        _check_not_supported()
        if self.dataset_shuffle is not None:
            self.train_dataloader_shuffle = self.dataset_shuffle
        self._init_generation_batch_params()
        self.remove_unused_columns = False
        logger.info(f'Setting args.remove_unused_columns: {self.remove_unused_columns}')
        if self.truncation_strategy is None:
            self.truncation_strategy = 'left'
        if self.truncation_strategy not in {'left', 'delete'}:
            raise ValueError("GRPO requires `truncation_strategy 'left' or 'delete'`, "
                             f"Current value: `truncation_strategy='{self.truncation_strategy}'`.")
        # disable normalization, REAL https://arxiv.org/abs/2602.05630
        if self.loss_type == 'real':
            self.scale_rewards = 'none'
            logger.warning(
                f"[REAL] scale_rewards='{self.scale_rewards}' is ignored. "
                "It will be forced to 'none' because 'loss_type = real' does not support reward normalization.")

        if self.beta is None:
            self.beta = 0.04  # https://arxiv.org/abs/2402.03300
        if self.async_generate:
            logger.info('Using async mode. This is a approximate version which '
                        'will use the old weights to generate responses to accelerate. '
                        'This will ignore the `CLIP` of advantages, if you found the training '
                        'is unstable, you may consider using --async_generate false.')
        if 'soft_overlong' in self.reward_funcs:
            assert self.soft_cache_length is not None, \
                'The soft_cache_length must be set when using soft overlong rewards.'
            if self.soft_max_length is None:
                self.soft_max_length = self.max_completion_length
                logger.info(f'Auto-configured soft_max_length = max_completion_length {self.max_completion_length}')
        if not self.use_ray:
            assert self.use_vllm, 'use_vllm must be True for Megatron GRPO'

        # Mirror deploy_args: gym_env implies use_gym_env unless the user said otherwise.
        if self.use_gym_env is None and self.gym_env is not None:
            self.use_gym_env = True


@dataclass
class MegatronTunerMixin:
    tuner_type: Literal['lora', 'full', 'lora_llm'] = 'full'
    freeze_llm: bool = False
    freeze_vit: bool = True
    freeze_aligner: bool = True
    # full
    freeze_parameters: List[str] = field(default_factory=list)
    freeze_parameters_regex: Optional[str] = None
    freeze_parameters_ratio: float = 0.  # 0 ~ 1
    trainable_parameters: List[str] = field(default_factory=list)
    trainable_parameters_regex: Optional[str] = None
    # lora
    target_modules: List[str] = field(default_factory=lambda: ['all-linear'])
    target_regex: Optional[str] = None
    modules_to_save: List[str] = field(default_factory=list)

    # lora
    lora_rank: int = 8
    lora_alpha: int = 32
    lora_dropout: float = 0.05
    lora_bias: Literal['none', 'all'] = 'none'
    lora_dtype: Literal['float16', 'bfloat16', 'float32', None] = None
    use_rslora: bool = False

    def __post_init__(self):
        if 0 < self.freeze_parameters_ratio < 1 and self.pipeline_model_parallel_size > 1:
            raise ValueError('`freeze_parameters_ratio` is not supported when `pipeline_model_parallel_size` > 1')
        if self.target_regex:
            self.target_modules = self.target_regex


@dataclass
class MegatronArguments(RLHFMegatronArgumentsMixin, MegatronTunerMixin):
    # training
    micro_batch_size: int = 1
    global_batch_size: int = 16
    recompute_granularity: Literal['selective', 'full', 'none'] = 'selective'
    recompute_method: Literal['uniform', 'block'] = None
    recompute_num_layers: Optional[int] = None
    recompute_modules: List[str] = field(default_factory=lambda: ['core_attn'])
    train_iters: Optional[int] = None
    num_train_epochs: Optional[int] = None

    masked_softmax_fusion: bool = True
    bias_dropout_fusion: bool = True
    bias_activation_fusion: bool = True
    apply_rope_fusion: bool = False
    gradient_accumulation_fusion: bool = True
    cross_entropy_loss_fusion: bool = True
    cross_entropy_fusion_impl: Literal['native', 'te'] = 'native'
    calculate_per_token_loss: Optional[bool] = None
    attention_backend: str = 'flash'  # flash, fused, unfused, local, auto
    optimizer: Literal['adam', 'sgd', 'muon', 'dist_muon'] = 'adam'
    optimizer_cpu_offload: bool = False
    optimizer_offload_fraction: float = 1.
    optimizer_cuda_graph: bool = False
    use_precision_aware_optimizer: bool = False
    main_grads_dtype: Literal['fp32', 'bf16'] = 'fp32'
    main_params_dtype: Literal['fp32', 'fp16'] = 'fp32'
    exp_avg_dtype: Literal['fp32', 'fp16', 'bf16', 'fp8'] = 'fp32'
    exp_avg_sq_dtype: Literal['fp32', 'fp16', 'bf16', 'fp8'] = 'fp32'
    manual_gc: bool = False
    manual_gc_steps: int = 0
    manual_gc_eval: bool = True

    # data
    seed: int = 42
    train_dataloader_shuffle: bool = True
    dataloader_num_workers: int = 4
    dataloader_pin_memory: bool = True
    dataloader_persistent_workers: bool = False
    dataloader_prefetch_factor: int = 2
    data_sharding: bool = False
    group_by_length: bool = False
    te_rng_tracker: bool = False
    data_parallel_random_init: Optional[bool] = False
    padding_free: bool = True
    mlp_padding_free: bool = False

    # learning rate
    lr_warmup_init: float = 0.
    lr: Optional[float] = None
    lr_decay_style: Literal['constant', 'linear', 'cosine', 'inverse-square-root', 'WSD'] = 'cosine'
    # The default is None, which will be set to `train_iters`.
    lr_decay_iters: Optional[int] = None
    lr_warmup_iters: int = 0
    lr_warmup_fraction: Optional[float] = None
    min_lr: float = 0
    # wsd
    lr_wsd_decay_style: Literal['exponential', 'linear', 'cosine', 'minus_sqrt'] = 'exponential'
    lr_wsd_decay_iters: Optional[int] = None

    # regularization
    weight_decay: float = 0.1
    weight_decay_incr_style: Literal['constant', 'linear', 'cosine'] = 'constant'
    start_weight_decay: Optional[float] = None
    end_weight_decay: Optional[float] = None
    clip_grad: float = 1.
    adam_beta1: float = 0.9
    adam_beta2: float = 0.95
    adam_eps: float = 1e-8
    sgd_momentum: float = 0.9
    muon_momentum: float = 0.9
    muon_split_qkv: bool = True
    muon_use_nesterov: bool = False
    muon_scale_mode: Literal['spectral', 'unit_rms_norm', 'shape_scaling'] = 'spectral'
    muon_fp32_matmul_prec: Literal['low', 'medium', 'high'] = 'medium'
    muon_coefficient_type: str = 'quintic'
    muon_num_ns_steps: int = 5
    muon_tp_mode: Literal['blockwise', 'duplicated', 'distributed'] = 'blockwise'
    muon_extra_scale_factor: float = 1.
    muon_scalar_optimizer: str = 'adam'

    # checkpoint
    output_dir: Optional[str] = None
    save_steps: int = 500
    no_save_optim: bool = False
    no_save_rng: bool = False
    mcore_model: Optional[str] = None
    mcore_adapter: Optional[str] = None
    no_load_optim: bool = False
    no_load_rng: bool = False
    finetune: bool = True
    perform_initialization: bool = False
    use_cpu_initialization: bool = False
    async_save: bool = False
    save_total_limit: Optional[int] = None
    metric_for_best_model: Optional[str] = None
    greater_is_better: Optional[bool] = None

    use_persistent_ckpt_worker: bool = False
    dist_ckpt_save_pre_mcore_014: bool = False
    dist_ckpt_optim_fully_reshardable: bool = False
    distrib_optim_fully_reshardable_mem_efficient: bool = False

    # dist
    local_rank: Optional[int] = None  # Compatible with DeepSpeed launch
    ddp_timeout: int = 18000000
    ddp_backend: Literal['nccl', 'gloo'] = 'nccl'
    use_distributed_optimizer: bool = True
    tensor_model_parallel_size: int = 1
    pipeline_model_parallel_size: int = 1
    decoder_first_pipeline_num_layers: Optional[int] = None
    decoder_last_pipeline_num_layers: Optional[int] = None
    account_for_embedding_in_pipeline_split: bool = False
    account_for_loss_in_pipeline_split: bool = False
    overlap_p2p_comm: bool = True
    batch_p2p_comm: Optional[bool] = None
    align_param_gather: bool = True

    sequence_parallel: bool = False
    context_parallel_size: int = 1
    cp_partition_mode: Literal['zigzag', 'contiguous'] = 'zigzag'
    sequence_packing_scheduler: Optional[Literal['dp_balanced', 'default_dynamic_cp']] = None
    tp_comm_overlap: bool = False
    overlap_grad_reduce: bool = False
    overlap_param_gather: bool = False
    overlap_param_gather_with_optimizer_step: bool = False
    align_grad_reduce: bool = True
    # Eagerly create NCCL communicators before the training loop to avoid the lazy
    # first-use allocation hitting the iteration-1 memory peak (Failed to CUDA calloc async).
    nccl_comm_warmup: bool = False
    virtual_pipeline_model_parallel_size: Optional[int] = None
    microbatch_group_size_per_vp_stage: Optional[int] = None
    pipeline_model_parallel_layout: Optional[str] = None
    expert_model_parallel_size: int = 1
    expert_tensor_parallel_size: int = 1

    # 'wandb', 'swanlab', 'tensorboard'
    report_to: List[str] = field(default_factory=lambda: ['tensorboard'])
    logging_steps: int = 5
    tensorboard_dir: Optional[str] = None
    tensorboard_queue_size: int = 50
    wandb_project: str = 'megatron-swift'
    wandb_exp_name: Optional[str] = None
    swanlab_project: str = 'megatron-swift'
    swanlab_exp_name: Optional[str] = None

    # evaluate
    eval_iters: int = -1
    eval_steps: Optional[int] = None

    # fp8
    fp8_format: Literal['e4m3', 'hybrid'] = None
    fp8_recipe: Literal['tensorwise', 'delayed', 'mxfp8', 'blockwise'] = 'delayed'
    fp8_param_gather: bool = False
    fp8_amax_history_len: int = 1024
    fp8_amax_compute_algo: Literal['most_recent', 'max'] = 'max'

    # fp4
    fp4_format: Literal['e2m1'] = None
    fp4_recipe: Literal['nvfp4'] = 'nvfp4'
    fp4_param_gather: bool = False

    # mixed precision
    fp16: Optional[bool] = None
    bf16: Optional[bool] = None
    apply_query_key_layer_scaling: Optional[bool] = None
    attention_softmax_in_fp32: bool = True
    accumulate_allreduce_grads_in_fp32: bool = False

    # moe
    moe_router_load_balancing_type: Optional[List[str]] = None
    moe_router_dtype: Literal['none', 'fp32', 'fp64'] = 'fp32'
    moe_token_dispatcher_type: Literal['allgather', 'alltoall', 'flex'] = 'alltoall'
    moe_enable_deepep: bool = False
    moe_grouped_gemm: bool = True
    moe_permute_fusion: bool = False
    moe_aux_loss_coeff: List[float] = 0.
    moe_z_loss_coeff: Optional[float] = None
    moe_shared_expert_overlap: bool = False
    moe_layer_recompute: bool = False  # compat mcore 0.12
    moe_expert_capacity_factor: Optional[float] = None
    moe_pad_expert_input_to_capacity: bool = False
    moe_token_drop_policy: Literal['probs', 'position'] = 'probs'

    # mtp
    mtp_num_layers: Optional[int] = None
    mtp_loss_scaling_factor: float = 0.1
    mtp_decoder_input_detach: bool = False
    mtp_shared_weights: bool = False

    # mcore-bridge / megatron-bridge
    bridge_backend: Literal['mcore-bridge', 'megatron-bridge'] = 'mcore-bridge'
    model: Optional[str] = None
    model_type: Optional[str] = None
    save_safetensors: bool = True
    adapters: List[str] = field(default_factory=list)
    ref_model: Optional[str] = None
    ref_adapters: List[str] = field(default_factory=list)
    use_hf: bool = False
    # None: use env var `MODELSCOPE_API_TOKEN`
    hub_token: Optional[str] = field(
        default=None, metadata={'help': 'SDK token can be found in https://modelscope.cn/my/myaccesstoken'})
    merge_lora: Optional[bool] = None
    max_shard_size: str = '5GB'

    # visual
    vit_gradient_checkpointing: Optional[bool] = None
    vit_gradient_checkpointing_kwargs: Optional[Union[dict, str]] = None
    vit_attn_impl: Optional[str] = None
    vit_lr: Optional[float] = None
    aligner_lr: Optional[float] = None

    # dsa
    dsa_indexer_loss_coeff: float = 0.
    dsa_indexer_use_sparse_loss: bool = False
    apply_dsa_kernel_fusion: bool = False
    # deepseek-v4
    csa_dense_mode: bool = False
    use_fused_mhc: bool = False
    mhc_recompute_layer_num: Optional[int] = None

    # other
    megatron_extra_kwargs: Optional[Union[dict, str]] = None
    language_model_only: bool = False
    check_model: bool = True
    torch_dtype: Optional[Union[torch.dtype, str]] = None
    rope_scaling: Optional[Union[dict, str]] = None
    apply_wd_to_qk_layernorm: bool = False
    linear_decoupled_in_proj: bool = False

    enable_dft_loss: bool = False
    enable_channel_loss: bool = False
    task_type: Literal['causal_lm', 'seq_cls', 'embedding', 'generative_reranker'] = None
    num_labels: Optional[int] = None
    problem_type: Literal['regression', 'single_label_classification', 'multi_label_classification'] = None
    # embedding (Matryoshka Representation Learning)
    # Dict[int, float], where the key is the embedding dimension and the value is the corresponding loss weight,
    # e.g. '{"32": 1.0, "64": 1.0, "128": 1.0}'.
    mrl_dims: Optional[Union[dict, str]] = None
    save_strategy: Literal['steps', 'epoch'] = 'steps'
    callbacks: List[str] = field(default_factory=list)

    @staticmethod
    def load_args_config(ckpt_dir: Optional[str]) -> Dict[str, Any]:
        res = {}
        if ckpt_dir is None:
            return res
        args_path = os.path.join(ckpt_dir, 'args.json')
        if os.path.exists(args_path):
            with open(args_path, 'r', encoding='utf-8') as f:
                old_args = json.load(f)
            keys = list(f.name for f in fields(MegatronTunerMixin))
            keys += ['mcore_model', 'task_type', 'num_labels', 'bridge_backend']
            for key in keys:
                old_value = old_args.get(key)
                if old_value is not None:
                    res[key] = old_value
            res.pop('mcore_adapter', None)
            if res['tuner_type'] == 'full':
                res.pop('mcore_model', None)
        return res

    def _set_default(self):
        if self.local_rank is None:
            self.local_rank = get_dist_setting()[1]
        if self.lr is None:
            if self.tuner_type == 'full':
                self.lr = 1e-5
            else:
                self.lr = 1e-4
        if self.task_type is None:
            self.task_type = 'causal_lm'
        if self.calculate_per_token_loss is None:
            self.calculate_per_token_loss = (self.task_type == 'causal_lm' and self.rlhf_type is None)

    def _init_mixed_precision(self):
        ModelArguments._init_mixed_precision(self)
        if self.apply_query_key_layer_scaling is None:
            self.apply_query_key_layer_scaling = self.fp16
        if self.apply_query_key_layer_scaling:
            os.environ['NVTE_APPLY_QK_LAYER_SCALING'] = '1'

    def _check_mcore_bridge(self):
        if self.language_model_only:
            require_version('mcore-bridge>=1.4.3', 'Please install "mcore-bridge>=1.4.3" to use language_model_only.')
            if self.tuner_type == 'lora_llm':
                raise ValueError('`tuner_type="lora_llm"` is not supported when `language_model_only=True`. '
                                 'Please use `tuner_type="lora"` instead.')

    def _check_bridge_backend(self):
        """Validate bridge_backend and associated constraints."""
        if self.bridge_backend == 'megatron-bridge':
            try:
                import megatron.bridge
            except ImportError:
                raise ImportError('bridge_backend="megatron-bridge" requires the `megatron-bridge` package. '
                                  'Install it via `pip install megatron-bridge` or use bridge_backend="mcore-bridge".')
            if self.tuner_type != 'full':
                raise ValueError('LoRA training is not yet supported with bridge_backend="megatron-bridge". '
                                 'Please use bridge_backend="mcore-bridge" for LoRA, or set tuner_type="full".')
        else:
            require_version('mcore-bridge>=1.4.0', 'Please install mcore-bridge via `pip install mcore-bridge -U`')
            from swift.megatron.init import _patch_mcore_bridge
            _patch_mcore_bridge()

    def __post_init__(self):
        if self.tuner_type != 'full':
            require_version('peft>=0.15', 'Please install peft>=0.15 to use LoRA in Megatron-SWIFT.')
        RLHFMegatronArgumentsMixin.__post_init__(self)
        MegatronTunerMixin.__post_init__(self)
        os.environ.setdefault('CUDA_DEVICE_MAX_CONNECTIONS', '1')
        self._check_mcore_bridge()
        self._check_bridge_backend()
        if self.recompute_granularity == 'none':
            self.recompute_granularity = None
        if self.recompute_granularity == 'selective' and self.recompute_method is not None:
            raise ValueError('recompute method is not yet supported for selective recomputing granularity')

        if self.group_by_length and self.padding_free:
            raise ValueError('group_by_length is not compatible with padding_free.')
        self._set_default()
        self._init_vpp_size()
        if self.vit_gradient_checkpointing is None:
            self.vit_gradient_checkpointing = not self.freeze_vit
        if isinstance(self.report_to, str):
            self.report_to = [self.report_to]
        self.model_info, self.model_meta = get_model_info_meta(
            self.model, model_type=self.model_type, use_hf=self.use_hf, hub_token=self.hub_token)

        # Megatron has a model_type parameter with the same name, so we need to avoid conflicts.
        self.model_type = self.model_info.model_type
        self.model_dir = self.model_info.model_dir
        self.is_multimodal = self.model_meta.is_multimodal
        if self.bridge_backend == 'megatron-bridge':
            self.megatron_model_meta = None
            if self.is_multimodal:
                raise ValueError('Multimodal training is not yet supported with bridge_backend="megatron-bridge". '
                                 'Please use bridge_backend="mcore-bridge" for multimodal models.')
            if self.task_type not in (None, 'causal_lm'):
                raise ValueError(f'task_type={self.task_type!r} is not yet supported with '
                                 f'bridge_backend="megatron-bridge".')
        else:
            self.megatron_model_meta = get_model_meta(self._get_mcore_model_type(self.model_meta))
            if self.megatron_model_meta is None:
                raise ValueError(f'Model: {self.model} is not supported.')
        self._init_teacher_model()
        if self.apply_wd_to_qk_layernorm and self.model_type not in {'qwen3_next', 'qwen3_5', 'qwen3_5_moe'}:
            raise ValueError('apply_wd_to_qk_layernorm is only supported for qwen3_next, qwen3_5 and qwen3_5_moe')
        if self.pipeline_model_parallel_size == 1 and (self.decoder_first_pipeline_num_layers is not None
                                                       or self.decoder_last_pipeline_num_layers is not None):
            raise ValueError('pipeline_model_parallel_size must be greater than 1 if you want to set '
                             'decoder_first_pipeline_num_layers or decoder_last_pipeline_num_layers.')
        # compat megatron-core
        self.fp8 = self.fp8_format
        self.fp4 = self.fp4_format

        if self.megatron_extra_kwargs is not None:
            self.megatron_extra_kwargs = json_parse_to_dict(self.megatron_extra_kwargs)
        if self.mrl_dims is not None:
            self.mrl_dims = json_parse_to_dict(self.mrl_dims)
            self.mrl_dims = {int(k): float(v) for k, v in self.mrl_dims.items()}
        if self.task_type not in {'causal_lm', 'generative_reranker'}:
            self.untie_embeddings_and_output_weights = True
        if self.vit_gradient_checkpointing_kwargs is not None:
            self.vit_gradient_checkpointing_kwargs = json_parse_to_dict(self.vit_gradient_checkpointing_kwargs)
        if self.gradient_accumulation_fusion:
            try:
                import apex
            except ImportError:
                logger.warning('apex is not installed, so gradient accumulation fusion is disabled.')
                self.gradient_accumulation_fusion = False
        self.callbacks += ['print', 'default_flow']
        self.callbacks += self.report_to
        if self.save_total_limit is not None:
            if self.async_save:
                raise ValueError('async_save is not supported with save_total_limit.')
            if self.save_total_limit < 2:
                raise ValueError('save_total_limit must be greater than or equal to 2.')
        if self.metric_for_best_model is None:
            self.metric_for_best_model = 'reward' if self.rlhf_type == 'grpo' else 'loss'
        if self.greater_is_better is None and self.metric_for_best_model is not None:
            self.greater_is_better = 'loss' not in self.metric_for_best_model
        if isinstance(self.ref_adapters, str):
            self.ref_adapters = [self.ref_adapters]
        if self.eval_steps is None:
            self.eval_steps = self.save_steps
        if self.merge_lora is None:
            self.merge_lora = self.save_safetensors
        if self.tuner_type == 'lora_llm':
            if not self.is_multimodal:
                raise ValueError('`tuner_type="lora_llm"` is only supported for multimodal models.')
            if not self.merge_lora:
                raise ValueError('`merge_lora` must be True when using `--tuner_type lora_llm`')
            if not self.no_save_optim:
                raise ValueError('`no_save_optim` must be True when using `--tuner_type lora_llm`')
        if self.adapters or self.ref_adapters or self.mcore_adapter or self.mcore_ref_adapter:
            if self.tuner_type == 'full':
                self.tuner_type = 'lora'
                logger.info('Setting args.tuner_type: lora')
        if self.adapters:
            self._load_adapter_config()
        self._init_mixed_precision()
        self._init_multimodal_full()
        self._map_dtype()
        self._init_weigh_decay()
        self._init_attention_backend()
        if self.sequence_parallel and self.tensor_model_parallel_size <= 1:
            self.sequence_parallel = False
        if isinstance(self.moe_aux_loss_coeff, list) and len(self.moe_aux_loss_coeff) == 1:
            self.moe_aux_loss_coeff = self.moe_aux_loss_coeff[0]
        if isinstance(self.moe_router_load_balancing_type, list) and len(self.moe_router_load_balancing_type) == 1:
            self.moe_router_load_balancing_type = self.moe_router_load_balancing_type[0]
        if self.tp_comm_overlap and not self.sequence_parallel:
            raise ValueError('Tensor parallel communication/GEMM overlap can happen only when '
                             'sequence parallelism is enabled')

        self._init_distributed()
        self._check_muon()

    def _init_attention_backend(self):
        if self.attention_backend.startswith('flash_'):
            from transformer_engine.pytorch.attention.dot_product_attention.utils import FlashAttentionUtils as fa_utils

            fa_version = int(self.attention_backend[len('flash_'):])
            assert fa_version in (2, 3, 4), (f'Unsupported flash attention version: {fa_version}. '
                                             f'Supported: flash_2, flash_3, flash_4.')
            available = {2: fa_utils.is_installed, 3: fa_utils.v3_is_installed, 4: fa_utils.v4_is_installed}
            if not available[fa_version]:
                raise ValueError(f'flash-attn v{fa_version} is not installed. '
                                 f'Detected installations: FA2={available[2]}, FA3={available[3]}, FA4={available[4]}.')

            if fa_version != 2:
                fa_utils.is_installed = False
            if fa_version != 3:
                fa_utils.v3_is_installed = False
            if fa_version != 4:
                fa_utils.v4_is_installed = False
            logger.info(f'Forcing Flash Attention v{fa_version} as the attention backend.')
            self.attention_backend = 'flash'
        self.attention_backend = AttnBackend[self.attention_backend]

    def _init_distributed(self):
        initialize_megatron(self)
        total_model_size = (
            self.tensor_model_parallel_size * self.pipeline_model_parallel_size * self.context_parallel_size)
        # world_size is initialized in initialize_megatron
        self.data_parallel_size = self.world_size // total_model_size
        # Gradient Accumulation
        micro_batch_times_data_parallel_size = self.micro_batch_size * self.data_parallel_size
        if self.global_batch_size % micro_batch_times_data_parallel_size != 0:
            raise ValueError(f'global batch size ({self.global_batch_size}) is not divisible by micro batch size '
                             f'({self.micro_batch_size}) times data parallel size ({self.data_parallel_size}).')
        self.num_microbatches = self.global_batch_size // micro_batch_times_data_parallel_size
        if self.num_microbatches == 0:
            raise ValueError('global_batch_size must be >= `data_parallel_size * micro_batch_size` '
                             f'to have at least one micro-batch. global_batch_size: {self.global_batch_size}, '
                             f'data_parallel_size: {self.data_parallel_size}, '
                             f'micro_batch_size: {self.micro_batch_size}.')

    def _get_mcore_model_type(self, model_meta):
        model_type = model_meta.model_type
        mcore_model_type = self.model_meta.mcore_model_type
        if mcore_model_type is None:
            mcore_model_type = get_mcore_model_type(model_type)
        return mcore_model_type

    def _check_muon(self):
        # Code borrowed from NVIDIA/Megatron-LM
        if 'muon' in self.optimizer:
            if not mcore_016:
                raise ValueError('Muon optimizer requires "megatron-core>=0.16".')

            if self.optimizer == 'muon':
                assert not self.overlap_grad_reduce, (
                    'Muon optimizer does not support overlap grad reduce. Use dist_muon instead.')
                assert not self.overlap_param_gather, (
                    'Muon optimizer does not support overlap param gather. Use dist_muon instead.')
            # Muon optimizer does not support distributed optimizer for now.
            self.use_distributed_optimizer = False
            # compat mcore 0.17
            self.muon_nesterov = self.muon_use_nesterov

    def _init_teacher_model(self):
        if self.teacher_model is None:
            return
        self.teacher_model_info, self.teacher_model_meta = get_model_info_meta(
            self.teacher_model, model_type=self.teacher_model_type, use_hf=self.use_hf, hub_token=self.hub_token)
        self.teacher_model_type = self.teacher_model_info.model_type
        self.teacher_model_dir = self.teacher_model_info.model_dir
        if self.bridge_backend == 'megatron-bridge':
            self.teacher_megatron_model_meta = None
        else:
            self.teacher_megatron_model_meta = get_model_meta(self._get_mcore_model_type(self.teacher_model_meta))
            if self.teacher_megatron_model_meta is None:
                raise ValueError(f'Model: {self.teacher_model} is not supported.')

    def _init_vpp_size(self):
        if self.pipeline_model_parallel_layout is not None:
            from megatron.core.transformer.pipeline_parallel_layer_layout import PipelineParallelLayerLayout

            # Parse the input flattened layout to a list and get the vpp size.
            # We will validate the layout more carefully in the TransformerConfig constructor.
            num_stages = PipelineParallelLayerLayout.get_num_stages_from_str(self.pipeline_model_parallel_layout)
            assert num_stages % self.pipeline_model_parallel_size == 0, (
                f'The length of pipeline_model_parallel_layout must be divisible'
                f' by pipeline_model_parallel_size ({num_stages=},'
                f' {self.pipeline_model_parallel_size=})')
            self.virtual_pipeline_model_parallel_size = num_stages // self.pipeline_model_parallel_size
        elif self.virtual_pipeline_model_parallel_size is not None:
            self.virtual_pipeline_model_parallel_size = self.virtual_pipeline_model_parallel_size
        if self.virtual_pipeline_model_parallel_size == 1:
            self.virtual_pipeline_model_parallel_size = None
        if self.virtual_pipeline_model_parallel_size is None:
            self.overlap_p2p_comm = False
            self.align_param_gather = False
        if self.batch_p2p_comm is None:
            self.batch_p2p_comm = not self.overlap_p2p_comm

    def _load_adapter_config(self):
        assert len(self.adapters) == 1, 'Currently only support one adapter'
        adapter_path = self.adapters[0]
        adapter_config_path = os.path.join(adapter_path, 'adapter_config.json')
        adapter_config = {}
        if os.path.exists(adapter_config_path):
            with open(adapter_config_path, 'r', encoding='utf-8') as f:
                adapter_config = json.load(f)
        mapping = {'r': 'lora_rank', 'bias': 'lora_bias'}
        for k in ['lora_alpha', 'lora_dropout', 'use_rslora']:
            mapping[k] = k
        for k, v in adapter_config.items():
            if k not in mapping:
                continue
            k = mapping[k]
            if v != getattr(self, k):
                setattr(self, k, v)
                logger.info(f'Setting {k}: {v}')

    def init_iters(self, train_dataset, val_dataset):
        data_parallel_size = mpu.get_data_parallel_world_size()
        step_batch_size = self.micro_batch_size * data_parallel_size
        num_generations = self.num_generations if self.rlhf_type == 'grpo' else 1
        # TODO: Check if it causes duplicate saving at the end.
        if self.save_strategy == 'epoch':
            if hasattr(train_dataset, '__len__'):
                dataset_sample = len(train_dataset) // step_batch_size * step_batch_size * num_generations
                self.save_steps = dataset_sample // self.global_batch_size
                self.eval_steps = self.save_steps
            else:
                raise ValueError('streaming dataset is not supported with `--save_strategy epoch`.')
        if self.num_train_epochs is not None:
            if hasattr(train_dataset, '__len__'):
                dataset_sample = len(train_dataset) // step_batch_size * step_batch_size * num_generations
                self.train_iters = dataset_sample * self.num_train_epochs // self.global_batch_size
            elif self.train_iters is None:
                raise ValueError(
                    'You are using a streaming training dataset. Please explicitly specify `--train_iters`.')
        if self.eval_iters < 0:
            if val_dataset is None:
                self.eval_iters = 0
            elif hasattr(val_dataset, '__len__'):
                dataset_sample = len(val_dataset) // step_batch_size * step_batch_size
                dataset_sample = dataset_sample * num_generations
                self.eval_iters = max(dataset_sample // self.global_batch_size, 1)
            else:
                raise ValueError(
                    'You are using a streaming validation dataset. Please explicitly specify `--eval_iters`.')
            logger.info(f'Setting args.eval_iters: {self.eval_iters}')

        data_parallel_size = mpu.get_data_parallel_world_size()
        step_batch_size = self.micro_batch_size * data_parallel_size
        # To avoid errors caused by the validation set being insufficient to complete a single step.
        if val_dataset is not None and hasattr(val_dataset, '__len__') and len(val_dataset) < step_batch_size:
            val_dataset = None
        if val_dataset is None:
            self.eval_iters = 0

    def _init_multimodal_full(self):
        if not self.is_multimodal:
            return
        visual_cls = self.megatron_model_meta.visual_cls
        if self.tuner_type == 'full' and self.is_multimodal and visual_cls is not None and not self.language_model_only:
            vision_tower = [f'visual.{vit}' for vit in getattr(visual_cls, '_vision_tower', [])]
            aligner = [f'visual.{aligner}' for aligner in getattr(visual_cls, '_aligner', [])]
            generator = [f'visual.{generator}' for generator in getattr(visual_cls, '_generator', [])]
            if self.freeze_llm:
                self.freeze_parameters.append('language_model')
            if self.freeze_vit:
                self.freeze_parameters += vision_tower
            if self.freeze_aligner:
                self.freeze_parameters += aligner
            else:
                self.trainable_parameters += aligner
            self.freeze_parameters += generator
            if self.freeze_parameters:
                logger.info(f'freeze_parameters: {self.freeze_parameters}')
            if self.trainable_parameters:
                logger.info(f'additional trainable_parameters: {self.trainable_parameters}')

    def _map_dtype(self):
        dtype_map = {
            'fp32': torch.float32,
            'bf16': torch.bfloat16,
            'fp16': torch.float16,
            'fp8': torch.uint8,
        }
        self.main_grads_dtype = dtype_map[self.main_grads_dtype]
        self.main_params_dtype = dtype_map[self.main_params_dtype]
        self.exp_avg_dtype = dtype_map[self.exp_avg_dtype]
        self.exp_avg_sq_dtype = dtype_map[self.exp_avg_sq_dtype]
        if self.fp16:
            self.torch_dtype = torch.float16
        elif self.bf16:
            self.torch_dtype = torch.bfloat16
            if self.main_grads_dtype == torch.float32:
                self.accumulate_allreduce_grads_in_fp32 = True

    def _init_weigh_decay(self):
        if self.weight_decay_incr_style == 'constant':
            assert self.start_weight_decay is None
            assert self.end_weight_decay is None
            self.start_weight_decay = self.end_weight_decay = self.weight_decay
        else:
            assert self.start_weight_decay is not None
            assert self.end_weight_decay is not None
