Create README.md

Updated project name and description in README.

README.md

@@ -1,55 +1,61 @@
-# LLM-Dolly
+**Read this in other languages: [English](README.md), [中文](README_zh.md).**
-LLM-Dolly is a custom LLM built from scratch   
+# LLM built from scratch  
-持续优化更新中...
+Continuously optimizing and updating...
 
-🗓️ 更新日志
+🗓️ Changelog
 
-#### 2025-06-05
+#### 2025-06-22
-- 📝 上传tokenizer的训练权重（1.56字符/token），调整tokenizer的格式和训练方式，和qwen2的tokenizer风格保持一致.
+- 📝 Added reward model training, PPO training, GRPO training
 
 <details>
-<summary>查看更多</summary>
+<summary>Click for more details</summary>
+  
+#### 2025-06-19
+- 📝 Pretokenized corpus and updated cold start code.
+
+#### 2025-06-05
+- 📝 Uploaded tokenizer training weights (1.56 chars/token), adjusted tokenizer format and training method to align with the style of the Qwen2 tokenizer.
 
 #### 2025-05-26
-- 📝 增加LLM文档，涵盖LLM全周期技术点。
+- 📝 Added LLM documentation covering full lifecycle technical points of LLMs.
 
 #### 2025-05-22
-- 📝 添加支持MoE模型，训练资源占用不稳定，测试模型Experts=8, experts_per_tok=4,GPU显存从60%-94%跳动,单个GPU利用率0%-100%跳动。当调小batch_size后，显存占用稳定在90%，单个GPU利用率稳定在90%以上，偶现30%的利用率。
+- 📝 Added support for MoE models. Training resource usage is unstable. Tested model with Experts=8, experts_per_tok=4. GPU memory fluctuated from 60% to 94%, and single GPU utilization fluctuated from 0% to 100%. After reducing the batch_size, memory usage stabilized at around 90%, single GPU utilization stabilized above 90%, occasionally dropping to around 30%.
 
 #### 2025-05-20
-- 📝 添加支持jsonl文件训练.
+- 📝 Added support for training with jsonl files.
 
 #### 2025-05-19
-- 📝 添加支持使用deepspeed训练代码，测试训练中最大batch_size提升38%，训练速度提升9.6%.
+- 📝 Added support for training code using DeepSpeed. Testing showed a maximum batch_size increase of 38% and a training speed increase of 9.6%.
 
 #### 2025-05-16
-- 📝 添加dolly_llm的预训练代码，进行一次预训练测试：模型0.6B，语料500M，46GB*4显卡.
+- 📝 Added pre-training code for dolly_llm. Conducted a pre-training test: model 0.6B, corpus 500M, using 46GB * 4 GPUs.
 
 #### 2025-05-14
-- ✅ 将dolly_llm作为pip包，进行安装
+- ✅ Released dolly_llm as a pip package for installation.
 
 #### 2025-05-09
-- 📝 优化[RMSNorm](https://arxiv.org/pdf/1910.07467)、MLP、[RoPE](https://arxiv.org/pdf/2104.09864)代码。
+- 📝 Optimized code for [RMSNorm](https://arxiv.org/pdf/1910.07467), MLP, and [RoPE](https://arxiv.org/pdf/2104.09864).
 
 #### 2025-05-07
-- 📝 使用transformers格式规范modeling和configuration，并设计修改v0.1版modeling_dolly，参数量11.5B。
+- 📝 Standardized modeling and configuration using the transformers format, and designed/modified modeling_dolly v0.1 with 11.5B parameters.
 
 #### 2025-05-06
-- 📝 实现 configuration_dolly类,以及添加v0.0版modeling_dolly。
+- 📝 Implemented the configuration_dolly class and added modeling_dolly v0.0.
 
 #### 2025-04-30
-- 📝 添加Tokenizer的[BBPE](https://arxiv.org/pdf/1909.03341)方式训练。
+- 📝 Added [BBPE](https://arxiv.org/pdf/1909.03341) method training for the Tokenizer.
 
 #### 2025-04-29
-- 📝 添加tokenizer构建代码，支持sentencepiece和transfomers的[BPE](https://arxiv.org/pdf/1508.07909)，支持从文本构建和从已有的tokenzier构建。
+- 📝 Added tokenizer construction code, supporting sentencepiece and transformers' [BPE](https://arxiv.org/pdf/1508.07909). Supports building from text and from existing tokenizers.
 
 #### 2025-04-24
-- ✅ 测试从transformers构建自定义的LLM模型结构。
+- ✅ Tested constructing custom LLM model architectures from transformers.
 </details>
 
-## 致谢
+## Acknowledgments
 
-特别感谢以下资源和文章的帮助：
+Special thanks to the following resources and articles for their assistance:
 
 - [Huggingface transformers](https://github.com/huggingface/transformers)
 - [Qwen](https://huggingface.co/Qwen)

README_zh.md

@@ -0,0 +1,65 @@
+**Read this in other languages: [English](README.md), [中文](README_zh.md).**
+# LLM built from scratch   
+持续优化更新中...
+
+🗓️ 更新日志
+
+#### 2025-06-22
+- 📝 增加奖励模型训练、PPO训练、GRPO训练
+
+<details>
+<summary>查看更多</summary>
+  
+#### 2025-06-19
+- 📝 pretoken corpus and update cold start code.
+
+#### 2025-06-05
+- 📝 上传tokenizer的训练权重（1.56字符/token），调整tokenizer的格式和训练方式，和qwen2的tokenizer风格保持一致.
+
+#### 2025-05-26
+- 📝 增加LLM文档，涵盖LLM全周期技术点。
+
+#### 2025-05-22
+- 📝 添加支持MoE模型，训练资源占用不稳定，测试模型Experts=8, experts_per_tok=4,GPU显存从60%-94%跳动,单个GPU利用率0%-100%跳动。当调小batch_size后，显存占用稳定在90%，单个GPU利用率稳定在90%以上，偶现30%的利用率。
+
+#### 2025-05-20
+- 📝 添加支持jsonl文件训练.
+
+#### 2025-05-19
+- 📝 添加支持使用deepspeed训练代码，测试训练中最大batch_size提升38%，训练速度提升9.6%.
+
+#### 2025-05-16
+- 📝 添加dolly_llm的预训练代码，进行一次预训练测试：模型0.6B，语料500M，46GB*4显卡.
+
+#### 2025-05-14
+- ✅ 将dolly_llm作为pip包，进行安装
+
+#### 2025-05-09
+- 📝 优化[RMSNorm](https://arxiv.org/pdf/1910.07467)、MLP、[RoPE](https://arxiv.org/pdf/2104.09864)代码。
+
+#### 2025-05-07
+- 📝 使用transformers格式规范modeling和configuration，并设计修改v0.1版modeling_dolly，参数量11.5B。
+
+#### 2025-05-06
+- 📝 实现 configuration_dolly类,以及添加v0.0版modeling_dolly。
+
+#### 2025-04-30
+- 📝 添加Tokenizer的[BBPE](https://arxiv.org/pdf/1909.03341)方式训练。
+
+#### 2025-04-29
+- 📝 添加tokenizer构建代码，支持sentencepiece和transfomers的[BPE](https://arxiv.org/pdf/1508.07909)，支持从文本构建和从已有的tokenzier构建。
+
+#### 2025-04-24
+- ✅ 测试从transformers构建自定义的LLM模型结构。
+</details>
+
+## 致谢
+
+特别感谢以下资源和文章的帮助：
+
+- [Huggingface transformers](https://github.com/huggingface/transformers)
+- [Qwen](https://huggingface.co/Qwen)
+- [Neural Machine Translation of Rare Words with Subword Units](https://arxiv.org/pdf/1508.07909)
+- [Neural Machine Translation with Byte-Level Subwords](https://arxiv.org/pdf/1909.03341)
+- [Root Mean Square Layer Normalization](https://arxiv.org/pdf/1910.07467)
+- [RoPE RoFormer: Enhanced Transformer with Rotary Position Embedding](https://arxiv.org/pdf/2104.09864)

documents/tech.md

@@ -0,0 +1,56 @@
+一、RMSNorm操作
+代码实现几乎一样，但在返回值的操作上略有不同：
+OSS: 
+    return (self.weight * hidden_states).to(input_dtype) 
+Qwen3:
+    return self.weight * hidden_states.to(input_dtype) 
+更推荐OSS的方式
+
+二、MoE层
+1.专家网络权重形式不一样
+OSS采用大矩阵存储所有的MLP,且有bias，gate_up_proj形状为num_experts*hidden_size*(2*expert_dim),合并了两个上采样的权重；
+down_proj的权重形状为：num_experts*expert_dim*hidden_size
+
+Qwen3采用存储的粒度更细，每个专家的权重独立存储，且门控和上采样权重分开，不使用bias,形状均为hidden_size*expert_dim;
+下采样权重为：expert_dim*hiddent_size
+所有专家的权重定义为：nn.ModuleList(
+            [Qwen3MoeMLP(config, intermediate_size=config.moe_intermediate_size) for _ in range(self.num_experts)]
+        )
+
+2.专家的前向计算方式不一样
+Qwen3采用：Loop over all available experts
+OSS可以有两种方式：
+ When training is is more efficient to just loop over the experts and compute the output for each expert
+For inference we can sacrifice some memory and compute the output for all experts at once. 
+
+3.激活函数
+OSS使用：
+    glu = gate * torch.sigmoid(gate * self.alpha) 
+Qwen3使用：
+    silu
+4.MLP结构不一样，OSS including clamping and a residual connection
+OSS使用：
+                gate = gate.clamp(min=None, max=self.limit)
+                up = up.clamp(min=-self.limit, max=self.limit)
+                glu = gate * torch.sigmoid(gate * self.alpha)
+                gated_output = (up + 1) * glu
+                out = gated_output @ self.down_proj[expert_idx] + self.down_proj_bias[expert_idx]
+
+Qwen使用：
+        down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
+
+三、RoPE代码一样
+四、 eager_attention的计算
+    OSS使用了GQA和sink attention
+            combined_logits = torch.cat([attn_weights, sinks], dim=-1)
+            combined_logits = combined_logits - combined_logits.max(dim=-1, keepdim=True).values
+            probs = F.softmax(combined_logits, dim=-1, dtype=combined_logits.dtype)
+            scores = probs[..., :-1]
+    Qwen3使用了GQA，未使用sink attention
+            attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * scaling
+            attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype)
+五、 Attention类型不一样
+    OSS交替使用sliding_attention,full_attention
+    Qwen3支持sliding_window
+六、负载均衡
+    均使用Switch Transformer

src/step4_cold_start.py

@@ -0,0 +1,100 @@
+from transformers import AutoModelForCausalLM, AutoTokenizer
+from modeling_dolly import DollyForCausalLM
+from datasets import load_dataset
+from trl import SFTConfig, SFTTrainer, setup_chat_format
+import torch
+import os
+import deepspeed
+
+deepspeed.init_distributed()
+
+model_path = "./model"
+tokenizer_path = "./model" 
+
+model = DollyForCausalLM.from_pretrained(model_path)
+tokenizer = AutoTokenizer.from_pretrained(tokenizer_path)
+
+try:
+    model, tokenizer = setup_chat_format(model=model, tokenizer=tokenizer)
+except Exception as e:
+    print(f"Chat format setup failed: {e}, using default formatting")
+
+data_path = "./datas"
+train_data_files = [os.path.join(data_path, f) for f in os.listdir(data_path) if f.startswith("train")]
+valid_data_files = [os.path.join(data_path, f) for f in os.listdir(data_path) if f.startswith("valid")]
+
+file_ext = os.path.splitext(train_data_files[0])[1].lower() if train_data_files else '.json'
+data_type = {
+    '.json': 'json',
+    '.jsonl': 'json',
+    '.csv': 'csv',
+    '.txt': 'text'
+}.get(file_ext, 'json')
+
+ds = load_dataset(
+    data_type,
+    data_files={
+        'train': train_data_files,
+        'validation': valid_data_files
+    },
+    split=['train', 'validation']
+)
+
+finetune_name = "Dolly-SFT"
+sft_config = SFTConfig(
+    output_dir="./sft_output",
+    max_steps=1000000,
+    per_device_train_batch_size=16,  
+    per_device_eval_batch_size=2,
+    learning_rate=5e-5,
+    logging_steps=10,
+    save_steps=1000,
+    save_total_limit=2, 
+    evaluation_strategy="steps",
+    eval_steps=500,
+    load_best_model_at_end=True, 
+    resume_from_checkpoint=True,
+    hub_model_id=finetune_name,
+    deepspeed="./config/ds_config.json",  
+    fp16=True,  
+    gradient_accumulation_steps=8, 
+)
+
+trainer = SFTTrainer(
+    model=model,
+    args=sft_config,
+    train_dataset=ds[0],
+    eval_dataset=ds[1],
+    tokenizer=tokenizer,
+)
+
+checkpoint_dir = None
+if os.path.exists(sft_config.output_dir):
+    checkpoints = [d for d in os.listdir(sft_config.output_dir) if d.startswith("checkpoint")]
+    if checkpoints:
+        latest_checkpoint = max(
+            [os.path.join(sft_config.output_dir, d) for d in checkpoints],
+            key=os.path.getctime
+        )
+        checkpoint_dir = latest_checkpoint
+        print(f"Resuming training from checkpoint: {latest_checkpoint}")
+
+
+trainer.train(resume_from_checkpoint=checkpoint_dir)
+
+if trainer.is_world_process_zero():
+    trainer.save_model(f"./{finetune_name}")
+    
+    device = "cuda:0" if torch.cuda.is_available() else "cpu"
+    model = AutoModelForCausalLM.from_pretrained(f"./{finetune_name}").to(device)
+    tokenizer = AutoTokenizer.from_pretrained(f"./{finetune_name}")
+
+    prompt = "Write a haiku about programming"
+    messages = [{"role": "user", "content": prompt}]
+    formatted_prompt = tokenizer.apply_chat_template(messages, tokenize=False)
+
+    inputs = tokenizer(formatted_prompt, return_tensors="pt").to(device)
+    outputs = model.generate(**inputs, max_new_tokens=100)
+    print("\nAfter training:")
+    print(tokenizer.decode(outputs[0], skip_special_tokens=True))
+    #deepspeed --num_gpus 8 step4_cold_start.py

src/step5.0_train_reward_model.py

@@ -0,0 +1,125 @@
+import os
+import warnings
+import torch
+from datasets import load_from_disk
+from modeling_dolly import DollyForSequenceClassification
+from transformers import (
+    AutoModelForSequenceClassification,
+    AutoTokenizer,
+    HfArgumentParser,
+    TrainerCallback,
+    TrainingArguments
+)
+from trl import (
+    ModelConfig,
+    RewardConfig,
+    RewardTrainer,
+    ScriptArguments,
+    get_kbit_device_map,
+    get_peft_config,
+    get_quantization_config,
+    setup_chat_format,
+)
+from transformers.trainer_utils import get_last_checkpoint
+
+class ResumeTrainingCallback(TrainerCallback):
+    def on_train_begin(self, args, state, control, **kwargs):
+        if state.resume_from_checkpoint is None:
+            checkpoint = get_last_checkpoint(args.output_dir)
+            if checkpoint is not None:
+                control.should_resume = True
+
+if __name__ == "__main__":
+    parser = HfArgumentParser((ScriptArguments, RewardConfig, ModelConfig))
+    script_args, training_args, model_args = parser.parse_args_into_dataclasses()
+    training_args.gradient_checkpointing_kwargs = dict(use_reentrant=False)
+    
+    if script_args.deepspeed_config:
+        training_args.deepspeed = script_args.deepspeed_config
+
+    checkpoint_dir = get_last_checkpoint(training_args.output_dir)
+    if checkpoint_dir is not None:
+        training_args.resume_from_checkpoint = checkpoint_dir
+        print(f"恢复训练自检查点: {checkpoint_dir}")
+    torch_dtype = (
+        model_args.torch_dtype if model_args.torch_dtype in ["auto", None] else getattr(torch, model_args.torch_dtype)
+    )
+    quantization_config = get_quantization_config(model_args)
+    
+    model_kwargs = dict(
+        revision=model_args.model_revision,
+        device_map=get_kbit_device_map() if quantization_config is not None else None,
+        quantization_config=quantization_config,
+        use_cache=False if training_args.gradient_checkpointing else True,
+        torch_dtype=torch_dtype,
+    )
+    
+    tokenizer = AutoTokenizer.from_pretrained(
+        os.path.join("model", model_args.model_name_or_path),
+        trust_remote_code=model_args.trust_remote_code,
+        use_fast=True
+    )
+    model = DollyForSequenceClassification.from_pretrained(
+        os.path.join("model", model_args.model_name_or_path),
+        num_labels=1,
+        trust_remote_code=model_args.trust_remote_code,
+        **model_kwargs
+    )
+    
+    model.config.pad_token_id = tokenizer.pad_token_id
+    if tokenizer.chat_template is None:
+        model, tokenizer = setup_chat_format(model, tokenizer)
+
+    if model_args.use_peft and model_args.lora_task_type != "SEQ_CLS":
+        warnings.warn(
+            "You are using a `task_type` that is different than `SEQ_CLS` for PEFT. This will lead to silent bugs"
+            " Make sure to pass --lora_task_type SEQ_CLS when using this script with PEFT.",
+            UserWarning,
+        )
+
+    dataset_path = os.path.join("datas", script_args.dataset_name)
+    if not os.path.exists(dataset_path):
+        raise ValueError(f"数据集路径不存在: {dataset_path}")
+    
+    dataset = load_from_disk(dataset_path)
+    trainer = RewardTrainer(
+        model=model,
+        tokenizer=tokenizer,
+        args=training_args,
+        train_dataset=dataset[script_args.dataset_train_split],
+        eval_dataset=dataset[script_args.dataset_test_split] if training_args.eval_strategy != "no" else None,
+        peft_config=get_peft_config(model_args) if model_args.use_peft else None,
+    )
+    trainer.add_callback(ResumeTrainingCallback())
+    trainer.train(resume_from_checkpoint=training_args.resume_from_checkpoint)
+    if not model_args.use_peft:
+        trainer.save_model(training_args.output_dir)
+    else:
+        trainer.model.save_pretrained(training_args.output_dir)
+    
+    if training_args.eval_strategy != "no":
+        metrics = trainer.evaluate()
+        trainer.log_metrics("eval", metrics)
+        trainer.save_metrics("eval", metrics)
+    tokenizer.save_pretrained(training_args.output_dir)
+
+
+"""
+torchrun --nproc_per_node=8 \
+    reward_training.py \
+    --model_name_or_path qwen_local \
+    --dataset_name my_reward_data \
+    --output_dir ./reward_model_output \
+    --per_device_train_batch_size 16 \
+    --num_train_epochs 3 \
+    --gradient_accumulation_steps 2 \
+    --learning_rate 2e-5 \
+    --eval_strategy steps \
+    --eval_steps 200 \
+    --save_strategy steps \
+    --save_steps 500 \
+    --max_length 2048 \
+    --deepspeed_config configs/deepspeed_z3.json \
+    --bf16 \
+    --report_to tensorboard
+"""

src/step5.1_PPO.py

@@ -0,0 +1,152 @@
+import os
+import torch
+from accelerate import PartialState
+from datasets import load_from_disk
+from transformers.trainer_utils import get_last_checkpoint
+from modeling_dolly import DollyForSequenceClassification
+from modeling_dolly import DollyForCausalLM
+from transformers import (
+    AutoModelForCausalLM,
+    AutoModelForSequenceClassification,
+    AutoTokenizer,
+    HfArgumentParser,
+)
+
+from trl import (
+    ModelConfig,
+    PPOConfig,
+    PPOTrainer,
+    ScriptArguments,
+    get_kbit_device_map,
+    get_peft_config,
+    get_quantization_config,
+)
+from trl.trainer.utils import SIMPLE_CHAT_TEMPLATE
+
+if __name__ == "__main__":
+    parser = HfArgumentParser((ScriptArguments, PPOConfig, ModelConfig))
+    script_args, training_args, model_args = parser.parse_args_into_dataclasses()
+    if script_args.deepspeed_config:
+        training_args.deepspeed = script_args.deepspeed_config
+        checkpoint_dir = None
+    if os.path.exists(training_args.output_dir):
+        checkpoint_dir = get_last_checkpoint(training_args.output_dir)
+        if checkpoint_dir is not None:
+            training_args.resume_from_checkpoint = checkpoint_dir
+            print(f"恢复训练自检查点: {checkpoint_dir}")
+
+    torch_dtype = (
+        model_args.torch_dtype if model_args.torch_dtype in ["auto", None] else getattr(torch, model_args.torch_dtype)
+    )
+    quantization_config = get_quantization_config(model_args)
+    model_kwargs = dict(
+        revision=model_args.model_revision,
+        attn_implementation=model_args.attn_implementation,
+        torch_dtype=torch_dtype,
+        device_map=get_kbit_device_map() if quantization_config is not None else None,
+        quantization_config=quantization_config,
+    )
+
+    tokenizer = AutoTokenizer.from_pretrained(
+        os.path.join("model", model_args.model_name_or_path), 
+        padding_side="left", 
+        trust_remote_code=model_args.trust_remote_code
+    )
+    
+
+    if tokenizer.chat_template is None:
+        tokenizer.chat_template = SIMPLE_CHAT_TEMPLATE
+    
+    reward_model_path = os.path.join("model", "reward_model") 
+    value_model = DollyForSequenceClassification.from_pretrained(
+        reward_model_path, 
+        trust_remote_code=model_args.trust_remote_code, 
+        num_labels=1,
+        **model_kwargs
+    )
+    
+    reward_model = DollyForSequenceClassification.from_pretrained(
+        reward_model_path, 
+        trust_remote_code=model_args.trust_remote_code, 
+        num_labels=1,
+        **model_kwargs
+    )
+    
+    policy_path = os.path.join("model", "policy_model") 
+    policy = DollyForCausalLM.from_pretrained(
+        policy_path, 
+        trust_remote_code=model_args.trust_remote_code,
+        **model_kwargs
+    )
+
+    peft_config = get_peft_config(model_args)
+    ref_policy = None
+    if peft_config is None:
+        ref_policy = DollyForCausalLM.from_pretrained(
+            policy_path, 
+            trust_remote_code=model_args.trust_remote_code,
+            **model_kwargs
+        )
+
+    dataset_path = os.path.join("datas", script_args.dataset_name)
+    if not os.path.exists(dataset_path):
+        raise ValueError(f"数据集路径不存在: {dataset_path}")
+    
+    dataset = load_from_disk(dataset_path)
+    
+    eval_samples = 100
+    train_dataset = dataset.select(range(len(dataset) - eval_samples))
+    eval_dataset = dataset.select(range(len(dataset) - eval_samples, len(dataset)))
+    dataset_text_field = "prompt" 
+
+    def prepare_dataset(dataset, tokenizer):
+        def tokenize(element):
+            outputs = tokenizer(
+                element[dataset_text_field],
+                padding=False,
+            )
+            return {"input_ids": outputs["input_ids"]}
+        
+        return dataset.map(
+            tokenize,
+            batched=True,
+            remove_columns=dataset.column_names,
+            num_proc=training_args.dataset_num_proc,
+        )
+
+    with PartialState().local_main_process_first():
+        train_dataset = prepare_dataset(train_dataset, tokenizer)
+        eval_dataset = prepare_dataset(eval_dataset, tokenizer)
+
+    trainer = PPOTrainer(
+        args=training_args,
+        tokenizer=tokenizer, 
+        model=policy,
+        ref_model=ref_policy,
+        reward_model=reward_model,
+        value_model=value_model,
+        train_dataset=train_dataset,
+        eval_dataset=eval_dataset,
+        peft_config=peft_config,
+    )
+    
+    trainer.train(resume_from_checkpoint=training_args.resume_from_checkpoint)
+
+    trainer.save_model(training_args.output_dir)
+    
+    trainer.generate_completions()
+    """
+    torchrun --nproc_per_node=8 \
+    train_ppo.py \
+    --model_name_or_path policy_model \
+    --dataset_name my_ppo_dataset \
+    --output_dir ./ppo_output \
+    --per_device_train_batch_size 8 \
+    --gradient_accumulation_steps 2 \
+    --learning_rate 1e-5 \
+    --num_train_epochs 3 \
+    --deepspeed_config config/deepspeed_z3.json \
+    --dataset_text_field prompt \
+    --reward_model_path model/reward_model \
+    --sft_model_path model/policy_model
+    """

src/step5.2_GRPO.py

@@ -0,0 +1,117 @@
+import os
+import torch
+from datasets import load_from_disk
+from modeling_dolly import DollyForSequenceClassification
+from modeling_dolly import DollyForCausalLM
+from transformers import (
+    AutoModelForCausalLM,
+    AutoTokenizer,
+    TrainingArguments,
+    TrainerCallback
+)
+from trl import GRPOConfig, GRPOTrainer
+from transformers.trainer_utils import get_last_checkpoint
+
+class ResumeTrainingCallback(TrainerCallback):
+    def on_train_begin(self, args, state, control, **kwargs):
+        if state.resume_from_checkpoint is None:
+            checkpoint = get_last_checkpoint(args.output_dir)
+            if checkpoint is not None:
+                control.should_resume = True
+
+class RewardModel:
+    def __init__(self, model_path, device="cuda"):
+        self.device = device
+        self.model = DollyForSequenceClassification.from_pretrained(model_path).to(device)
+        self.tokenizer = AutoTokenizer.from_pretrained(model_path)
+        self.model.eval()
+    
+    def __call__(self, completions, **kwargs):
+        inputs = self.tokenizer(
+            completions, 
+            padding=True,
+            truncation=True,
+            max_length=2048,
+            return_tensors="pt"
+        ).to(self.device)
+        
+        with torch.no_grad():
+            outputs = self.model(**inputs)
+        return outputs.logits.squeeze(-1).cpu().tolist()
+
+def main():
+    training_args = GRPOConfig(
+        output_dir="./grpo_output",
+        per_device_train_batch_size=32,
+        per_device_eval_batch_size=8,
+        num_train_epochs=3,
+        gradient_accumulation_steps=2,
+        learning_rate=1e-5,
+        logging_steps=10,
+        save_steps=500,
+        eval_steps=200,
+        warmup_ratio=0.1,
+        weight_decay=0.01,
+        fp16=True,
+        deepspeed="configs/deepspeed_z3.json",  
+        report_to="tensorboard",
+        remove_unused_columns=False,
+    )
+    
+    checkpoint_dir = get_last_checkpoint(training_args.output_dir)
+    if checkpoint_dir is not None:
+        training_args.resume_from_checkpoint = checkpoint_dir
+        print(f"恢复训练自检查点: {checkpoint_dir}")
+
+    policy_model_path = os.path.join("model", "qwen_local")  
+    tokenizer = AutoTokenizer.from_pretrained(policy_model_path)
+    model = DollyForCausalLM.from_pretrained(
+        policy_model_path,
+        torch_dtype=torch.bfloat16,
+        device_map="auto"
+    )
+    
+    if tokenizer.pad_token is None:
+        tokenizer.pad_token = tokenizer.eos_token
+    
+    reward_model_path = "./reward_model_output"  
+    reward_model = RewardModel(reward_model_path)
+    
+    dataset_path = os.path.join("datas", "my_grpo_data")  
+    if not os.path.exists(dataset_path):
+        raise ValueError(f"数据集路径不存在: {dataset_path}")
+    
+    dataset = load_from_disk(dataset_path)
+    
+    trainer = GRPOTrainer(
+        model=model,
+        tokenizer=tokenizer,
+        args=training_args,
+        reward_funcs=[reward_model],  
+        train_dataset=dataset,
+    )
+    
+    trainer.add_callback(ResumeTrainingCallback())
+    
+    trainer.train(resume_from_checkpoint=training_args.resume_from_checkpoint)
+    
+    trainer.save_model(os.path.join(training_args.output_dir, "final_model"))
+    
+    print("GRPO训练完成!")
+
+if __name__ == "__main__":
+    main()
+
+"""
+torchrun --nproc_per_node=8 \
+    train_grpo.py \
+    --output_dir ./grpo_output \
+    --per_device_train_batch_size 16 \
+    --num_train_epochs 3 \
+    --gradient_accumulation_steps 2 \
+    --learning_rate 1e-5 \
+    --logging_steps 10 \
+    --save_steps 500 \
+    --eval_steps 200 \
+    --fp16
+"""