INNER CODE UNIT · Python

TrainingArguments

HIT-SCIR/Chinese-Mixtral-8x7B · train.py:58

class TrainingArguments(transformers.TrainingArguments):
    cache_dir: Optional[str] = field(default=None)
    optim: str = field(default="adamw_torch")
    model_max_length: int = field(
        default=512,
        metadata={"help": "Maximum sequence length. Sequences will be right padded (and possibly truncated)."},
    )
    mode: Literal["pt", "sft"] = field(default="pt")
    neftune_noise_alpha: Optional[float] = field(default=None)
    bpe_dropout_alpha: Optional[float] = field(default=None)


def safe_save_model_for_hf_trainer(trainer: transformers.Trainer, output_dir: str):
    """Collects the state dict and dump to disk."""
    state_dict = trainer.model.state_dict()
    if trainer.args.should_save:
        cpu_state_dict = {key: value.cpu() for key, value in state_dict.items()}
        del state_dict

View source record →