---
license: mit
arxiv: 2412.17743
tags:
- optimizer_states
---

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

Both [**YuLan-Mini**](https://huggingface.co/yulan-team/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.

| Stage     | Curriculum Phase | 4K Context                  | 28K Context | Optimizer | Inference Arch | LAMBADA `Acc` | GSM8K `Acc` | HumanEval `pass@1` |
|-----------|------------------|-----------------------------|-------------|-----------|----------------|-------------|-------------|--------------------|
| Stable    | 5                | [YuLan-Mini-Phase5](https://huggingface.co/yulan-team/YuLan-Mini-Phase5) |       | | `yulanmini` | 53.85 | 3.41 | 12.26 |
| Stable    | 10               | [YuLan-Mini-Phase10](https://huggingface.co/yulan-team/YuLan-Mini-Phase10) |       | | `yulanmini` | 55.00 | 9.57 | 15.95 |
| Stable    | 15               | [YuLan-Mini-Phase15](https://huggingface.co/yulan-team/YuLan-Mini-Phase15) |       | | `yulanmini` | 55.81 | 13.81 | 16.99 |
| Stable    | 20               | [YuLan-Mini-Phase20](https://huggingface.co/yulan-team/YuLan-Mini-Phase20) | | ✅ | `yulanmini` | 55.81 | 21.39 | 20.79 |
| Stable    | 25 (1T tokens)   | [YuLan-Mini-Before-Annealing](https://huggingface.co/yulan-team/YuLan-Mini-Before-Annealing) | | ✅ | `yulanmini` | 55.67 | 29.94 | 34.06 |
| | | | | | | | |
| Annealing | 26               | YuLan-Mini-4K  |             | | `llama`\* | 64.72 | 66.65 | 61.60 |
| Annealing | 27               |                             | [YuLan-Mini](https://huggingface.co/yulan-team/YuLan-Mini)  | | `llama`\* | 65.67 | 68.46 | 64.00 |

\*: For easier inference and deployment, we merged the re-parameterized added parameters and scaling factors into the final released models ([**YuLan-Mini**](https://huggingface.co/yulan-team/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](https://huggingface.co/yulan-team/YuLan-Mini-Datasets) 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](https://huggingface.co/yulan-team/YuLan-Mini-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](https://huggingface.co/yulan-team/YuLan-Mini) base model to train your own Instruct version.
4. **Training dynamics** research. You can use YuLan-Mini's [intermediate checkpoints](https://huggingface.co/collections/yulan-team/yulan-mini-676d214b24376739b00d95f3) to explore internal changes during the pre-training process.
5. **Synthesize** your own data. You can use YuLan-Mini's [data pipeline](https://github.com/RUC-GSAI/YuLan-Mini) 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:

```json
{
  "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:

```json
{
  "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.

```python
trainer.train(resume_from_checkpoint=training_args.resume_from_checkpoint)
```

We provide an internal [training framework](https://github.com/RUC-GSAI/YuLan-Mini/tree/main/pretrain) 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](http://aibox.ruc.edu.cn/).

## License

- The code in this repository is released under the [MIT License](./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](https://arxiv.org/abs/2412.17743):

```
@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}, 
}
```