yulan-team/YuLan-Mini-Before-Annealing

IMPORTANT NOTICE: THIS IS AN INTERMEDIATE CHECKPOINT, NOT THE FINAL MODEL

Both YuLan-Mini and YuLan-Mini-Intermediate-4K were trained starting from this checkpoint.

This version includes the optimizer, allowing you to resume training using the Hugging Face Trainer and DeepSpeed Universal Checkpoint.

For easier inference and deployment, we merged the re-parameterized added parameters and scaling factors into the final released models (YuLan-Mini and YuLan-Mini-Intermediate-4K), enabling it to run on the Llama architecture. However, these parameters are still retained in the intermediate checkpoints from the training process.

What you can do with these pre-training resources

1. Pre-train your own LLM. You can use our data and curriculum to train a model that's just as powerful as YuLan-Mini.
2. Perform your own learning rate annealing. During the annealing phase, YuLan-Mini's learning ability is at its peak. You can resume training from the checkpoint before annealing and use your own dataset for learning rate annealing.
3. Fine-tune the Instruct version of the LLM. You can use the YuLan-Mini base model to train your own Instruct version.
4. Training dynamics research. You can use YuLan-Mini's intermediate checkpoints to explore internal changes during the pre-training process.
5. Synthesize your own data. You can use YuLan-Mini's data pipeline to clean and generate your own dataset.

Continual Training Tutorial

Step 1: Modify the config.json

Due to the implementation of Hugging Face Trainer, certain parameters are stored in the config.json file and cannot be modified through the Trainer's command-line arguments. Therefore, you need to update these parameters in the config.json file first, particularly:

• save_steps: The frequency of saving intermediate checkpoints.
• train_batch_size: The batch size per GPU (equivalent to per_device_train_batch_size in the Trainer). We used a batch size of 1008 (approximately 4M tokens) during the stable training stage. Maintaining this same batch size is equally important for training effectiveness.

Below is an example of a properly configured config.json file:

{
  "best_metric": null,
  "best_model_checkpoint": null,
  "epoch": 0.0,
  "eval_steps": 500,
  "global_step": 0,
  "is_hyper_param_search": false,
  "is_local_process_zero": true,
  "is_world_process_zero": true,
  "log_history": [],
  "logging_steps": 3,
  "max_steps": 0,
  "num_input_tokens_seen": 0,
  "num_train_epochs": 0,
  "save_steps": 250,
  "stateful_callbacks": {
    "TrainerControl": {
      "args": {
        "should_epoch_stop": false,
        "should_evaluate": false,
        "should_log": false,
        "should_save": true,
        "should_training_stop": true
      },
      "attributes": {}
    }
  },
  "total_flos": 0,
  "train_batch_size": 3,
  "trial_name": null,
  "trial_params": null
}

Step 2: Enable Universal Checkpointing in the DeepSpeed Configuration

To ensure DeepSpeed Integration loads the Universal Checkpoint, you need to enable this feature in the DeepSpeed configuration JSON file.

Here is an example of a ZeRO2 configuration with Universal Checkpointing enabled:

{
  "bf16": {
    "enabled": "auto"
  },
  "zero_optimization": {
    "stage": 2,
    "allgather_partitions": true,
    "allgather_bucket_size": 8e8,
    "overlap_comm": true,
    "reduce_scatter": true,
    "reduce_bucket_size": 8e8,
    "contiguous_gradients": true
  },
  "gradient_accumulation_steps": "auto",
  "gradient_clipping": "auto",
  "steps_per_print": 16,
  "train_batch_size": "auto",
  "train_micro_batch_size_per_gpu": "auto",
  "wall_clock_breakdown": false,
  "dump_state": true,
  "optimizer": {
    "type": "AdamW",
    "params": {
      "lr": "auto",
      "betas": "auto",
      "eps": "auto",
      "weight_decay": "auto"
    }
  },
  "checkpoint": {
    "load_universal": true
  }
}

Step 3: Resume Training

When calling trainer.train, include the resume_from_checkpoint argument to load the distributed optimizer state from the Universal Checkpoint and resume training.

trainer.train(resume_from_checkpoint=training_args.resume_from_checkpoint)

We provide an internal training framework for your reference, but you are free to choose other frameworks.

---

The Team

YuLan-Mini is developed and maintained by AI Box, Renmin University of China.

License

• The code in this repository is released under the MIT License.
• Policies regarding the use of model weights, intermediate optimizer states, and training data will be announced in future updates.
• Limitations: Despite our efforts to mitigate safety concerns and encourage the generation of ethical and lawful text, the probabilistic nature of language models may still lead to unexpected outputs. For instance, responses might contain bias, discrimination, or other harmful content. Please refrain from disseminating such content. We are not liable for any consequences arising from the spread of harmful information.

Citation

If you find YuLan-Mini helpful for your research or development, please cite our technical report:

@misc{hu2024yulanmini,
      title={YuLan-Mini: An Open Data-efficient Language Model}, 
      author={Yiwen Hu and Huatong Song and Jia Deng and Jiapeng Wang and Jie Chen and Kun Zhou and Yutao Zhu and Jinhao Jiang and Zican Dong and Wayne Xin Zhao and Ji-Rong Wen},
      year={2024},
      eprint={2412.17743},
      archivePrefix={arXiv},
      primaryClass={cs.CL},
      url={https://arxiv.org/abs/2412.17743}, 
}