Merge 744f16d802 into a3c7b01cc0
CharlesCheung996
GitHub
commit
32d15a9932
|
|
@ -1,6 +1,9 @@
|
|||
.venv*/
|
||||
__pycache__/
|
||||
*.pyc
|
||||
*.py[cod]
|
||||
.DS_Store
|
||||
out
|
||||
out/
|
||||
checkpoints/
|
||||
.smoke-out/
|
||||
website/
|
||||
docs-minimind/
|
||||
docs-minimind/
|
||||
|
|
|
|||
13
README.md
13
README.md
|
|
@ -284,6 +284,19 @@ print(torch.cuda.is_available())
|
|||
若 `cuda` 不可用，也仍可根据自身设备选择 `CPU` 或 `MPS` 运行，但训练速度与兼容性会有非常大的差异。
|
||||
如需安装或更换 PyTorch 版本，可参考 [torch_stable](https://download.pytorch.org/whl/torch_stable.html) 与[链接](https://blog.csdn.net/weixin_45456738/article/details/141029610?ops_request_misc=&request_id=&biz_id=102&utm_term=%E5%AE%89%E8%A3%85torch&utm_medium=distribute.pc_search_result.none-task-blog-2~all~sobaiduweb~default-2-141029610.nonecase&spm=1018.2226.3001.4187)
|
||||
|
||||
#### Apple Silicon / MPS
|
||||
|
||||
在支持 Metal 的 Apple Silicon Mac 上，训练脚本会按 `CUDA → MPS → CPU` 的顺序自动选择设备，也可显式传入 `--device mps`。MPS 支持 `float32`、`float16` 和 `bfloat16`；建议较新的 Apple Silicon 优先尝试默认的 `bfloat16`，遇到算子兼容问题时改用 `float32`。
|
||||
|
||||
```bash
|
||||
cd trainer
|
||||
PYTORCH_ENABLE_MPS_FALLBACK=1 python train_pretrain.py \
|
||||
--device mps --dtype bfloat16 --num_workers 0 \
|
||||
--batch_size 8 --accumulation_steps 32
|
||||
```
|
||||
|
||||
`PYTORCH_ENABLE_MPS_FALLBACK=1` 会让尚未支持的 MPS 算子回退到 CPU，兼容性更好但可能降低速度。确认训练全程没有回退后可移除此环境变量。Mac 只有一个 MPS 设备，请直接使用 `python train_xxx.py`，不要使用面向 CUDA/NCCL 多卡训练的 `torchrun`。MPS 不使用 CUDA 的 pinned memory，保存检查点后也会自动释放 MPS 缓存。
|
||||
|
||||
</details>
|
||||
|
||||
### 1' 下载数据
|
||||
|
|
|
|||
|
|
@ -6,7 +6,7 @@ import torch
|
|||
from transformers import AutoTokenizer, AutoModelForCausalLM, TextStreamer
|
||||
from model.model_minimind import MiniMindConfig, MiniMindForCausalLM
|
||||
from model.model_lora import *
|
||||
from trainer.trainer_utils import setup_seed, get_model_params
|
||||
from trainer.trainer_utils import setup_seed, get_model_params, get_default_device
|
||||
warnings.filterwarnings('ignore')
|
||||
|
||||
def init_model(args):
|
||||
|
|
@ -45,7 +45,7 @@ def main():
|
|||
parser.add_argument('--open_thinking', default=0, type=int, help="是否开启自适应思考（0=否，1=是）")
|
||||
parser.add_argument('--historys', default=0, type=int, help="携带历史对话轮数（需为偶数，0表示不携带历史）")
|
||||
parser.add_argument('--show_speed', default=1, type=int, help="显示decode速度（tokens/s）")
|
||||
parser.add_argument('--device', default='cuda' if torch.cuda.is_available() else 'cpu', type=str, help="运行设备")
|
||||
parser.add_argument('--device', default=get_default_device(), type=str, help="运行设备 (cuda:0/mps/cpu)")
|
||||
args = parser.parse_args()
|
||||
|
||||
prompts = [
|
||||
|
|
@ -91,4 +91,4 @@ def main():
|
|||
print(f'\n[Speed]: {gen_tokens / (time.time() - st):.2f} tokens/s\n\n') if args.show_speed else print('\n\n')
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
main()
|
||||
|
|
|
|||
|
|
@ -0,0 +1,68 @@
|
|||
"""Short MiniMind MPS training benchmark; it does not save checkpoints."""
|
||||
|
||||
import argparse
|
||||
import os
|
||||
import sys
|
||||
import time
|
||||
|
||||
import torch
|
||||
|
||||
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
|
||||
|
||||
from model.model_minimind import MiniMindConfig, MiniMindForCausalLM
|
||||
from trainer.trainer_utils import get_autocast_context, get_grad_scaler
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument('--batch-size', type=int, default=8)
|
||||
parser.add_argument('--seq-len', type=int, default=340)
|
||||
parser.add_argument('--hidden-size', type=int, default=768)
|
||||
parser.add_argument('--num-hidden-layers', type=int, default=8)
|
||||
parser.add_argument('--dtype', choices=['float32', 'float16', 'bfloat16'], default='bfloat16')
|
||||
parser.add_argument('--warmup-steps', type=int, default=1)
|
||||
parser.add_argument('--steps', type=int, default=3)
|
||||
args = parser.parse_args()
|
||||
|
||||
if not torch.backends.mps.is_available():
|
||||
raise RuntimeError('MPS is not available')
|
||||
|
||||
config = MiniMindConfig(
|
||||
hidden_size=args.hidden_size,
|
||||
num_hidden_layers=args.num_hidden_layers,
|
||||
max_position_embeddings=max(args.seq_len, 512),
|
||||
)
|
||||
model = MiniMindForCausalLM(config).to('mps').train()
|
||||
optimizer = torch.optim.AdamW(model.parameters(), lr=5e-4)
|
||||
scaler = get_grad_scaler('mps', args.dtype)
|
||||
input_ids = torch.randint(
|
||||
0, config.vocab_size, (args.batch_size, args.seq_len), device='mps'
|
||||
)
|
||||
|
||||
durations = []
|
||||
total_steps = args.warmup_steps + args.steps
|
||||
for step in range(total_steps):
|
||||
torch.mps.synchronize()
|
||||
started = time.perf_counter()
|
||||
with get_autocast_context('mps', args.dtype):
|
||||
loss = model(input_ids, labels=input_ids).loss
|
||||
scaler.scale(loss).backward()
|
||||
scaler.unscale_(optimizer)
|
||||
torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
|
||||
scaler.step(optimizer)
|
||||
scaler.update()
|
||||
optimizer.zero_grad(set_to_none=True)
|
||||
torch.mps.synchronize()
|
||||
elapsed = time.perf_counter() - started
|
||||
if step >= args.warmup_steps:
|
||||
durations.append(elapsed)
|
||||
|
||||
mean_seconds = sum(durations) / len(durations)
|
||||
tokens_per_second = args.batch_size * args.seq_len / mean_seconds
|
||||
print(f'mean_step_seconds={mean_seconds:.3f}')
|
||||
print(f'tokens_per_second={tokens_per_second:.1f}')
|
||||
print(f'loss={loss.detach().float().item():.4f}')
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
main()
|
||||
|
|
@ -0,0 +1,8 @@
|
|||
{"text":"苹果芯片可以通过 MPS 加速 PyTorch 模型训练。"}
|
||||
{"text":"MiniMind 是一个用于学习语言模型训练流程的小型项目。"}
|
||||
{"text":"梯度下降通过反向传播更新神经网络参数。"}
|
||||
{"text":"混合精度训练可以降低内存占用并提高吞吐量。"}
|
||||
{"text":"统一内存允许 CPU 和 GPU 共享同一块物理内存。"}
|
||||
{"text":"检查点可用于恢复模型参数和优化器状态。"}
|
||||
{"text":"注意力机制根据上下文计算不同词元之间的关系。"}
|
||||
{"text":"小规模冒烟测试可以快速发现训练环境问题。"}
|
||||
|
|
@ -0,0 +1,4 @@
|
|||
{"conversations":[{"role":"user","content":"MPS 是什么？","reasoning_content":"","tools":"","tool_calls":""},{"role":"assistant","content":"MPS 是 Apple 平台上的 GPU 加速框架。","reasoning_content":"","tools":"","tool_calls":""}]}
|
||||
{"conversations":[{"role":"user","content":"MiniMind 有什么用途？","reasoning_content":"","tools":"","tool_calls":""},{"role":"assistant","content":"它用于学习和实验小型语言模型训练流程。","reasoning_content":"","tools":"","tool_calls":""}]}
|
||||
{"conversations":[{"role":"user","content":"什么是梯度？","reasoning_content":"","tools":"","tool_calls":""},{"role":"assistant","content":"梯度描述损失函数相对模型参数的变化方向。","reasoning_content":"","tools":"","tool_calls":""}]}
|
||||
{"conversations":[{"role":"user","content":"为什么使用混合精度？","reasoning_content":"","tools":"","tool_calls":""},{"role":"assistant","content":"它通常能减少内存占用并提高训练速度。","reasoning_content":"","tools":"","tool_calls":""}]}
|
||||
|
|
@ -0,0 +1,98 @@
|
|||
import os
|
||||
import unittest
|
||||
from unittest import mock
|
||||
|
||||
import torch
|
||||
|
||||
from model.model_minimind import MiniMindConfig, MiniMindForCausalLM
|
||||
from trainer.trainer_utils import (
|
||||
empty_device_cache,
|
||||
get_autocast_context,
|
||||
get_default_device,
|
||||
get_grad_scaler,
|
||||
init_model,
|
||||
init_distributed_mode,
|
||||
LMForRewardModel,
|
||||
should_pin_memory,
|
||||
)
|
||||
|
||||
|
||||
class DeviceUtilsTest(unittest.TestCase):
|
||||
def test_default_device_prefers_mps_over_cpu(self):
|
||||
with mock.patch('torch.cuda.is_available', return_value=False), \
|
||||
mock.patch('torch.backends.mps.is_available', return_value=True):
|
||||
self.assertEqual(get_default_device(), 'mps')
|
||||
|
||||
def test_pin_memory_is_cuda_only(self):
|
||||
self.assertTrue(should_pin_memory('cuda:0'))
|
||||
self.assertFalse(should_pin_memory('mps'))
|
||||
self.assertFalse(should_pin_memory('cpu'))
|
||||
|
||||
def test_distributed_mps_fails_with_actionable_message(self):
|
||||
with mock.patch.dict(os.environ, {'RANK': '0'}):
|
||||
with self.assertRaisesRegex(RuntimeError, 'single process'):
|
||||
init_distributed_mode('mps')
|
||||
|
||||
@unittest.skipUnless(torch.backends.mps.is_available(), 'MPS is not available')
|
||||
def test_init_model_uses_default_device_when_device_is_omitted(self):
|
||||
config = MiniMindConfig(
|
||||
hidden_size=64,
|
||||
num_hidden_layers=2,
|
||||
max_position_embeddings=64,
|
||||
)
|
||||
model, _ = init_model(
|
||||
config,
|
||||
from_weight='none',
|
||||
tokenizer_path='model',
|
||||
save_dir='out',
|
||||
device=None,
|
||||
)
|
||||
|
||||
self.assertEqual(next(model.parameters()).device.type, 'mps')
|
||||
|
||||
def test_reward_model_uses_default_device_when_device_is_omitted(self):
|
||||
fake_model = mock.Mock()
|
||||
fake_model.to.return_value = fake_model
|
||||
fake_model.eval.return_value = fake_model
|
||||
|
||||
with mock.patch(
|
||||
'trainer.trainer_utils.get_default_device', return_value='mps'
|
||||
), mock.patch(
|
||||
'trainer.trainer_utils.AutoTokenizer.from_pretrained'
|
||||
), mock.patch(
|
||||
'trainer.trainer_utils.AutoModel.from_pretrained', return_value=fake_model
|
||||
):
|
||||
LMForRewardModel('unused', device=None, dtype=torch.float32)
|
||||
|
||||
fake_model.to.assert_called_once_with('mps')
|
||||
|
||||
@unittest.skipUnless(torch.backends.mps.is_available(), 'MPS is not available')
|
||||
def test_minimind_optimizer_step_on_mps_for_all_dtypes(self):
|
||||
for dtype in ('float32', 'float16', 'bfloat16'):
|
||||
with self.subTest(dtype=dtype):
|
||||
config = MiniMindConfig(
|
||||
hidden_size=64,
|
||||
num_hidden_layers=2,
|
||||
max_position_embeddings=64,
|
||||
)
|
||||
model = MiniMindForCausalLM(config).to('mps').train()
|
||||
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4)
|
||||
scaler = get_grad_scaler('mps', dtype)
|
||||
input_ids = torch.randint(0, config.vocab_size, (2, 32), device='mps')
|
||||
|
||||
with get_autocast_context('mps', dtype):
|
||||
loss = model(input_ids, labels=input_ids).loss
|
||||
|
||||
self.assertTrue(torch.isfinite(loss).item())
|
||||
scaler.scale(loss).backward()
|
||||
scaler.unscale_(optimizer)
|
||||
torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
|
||||
scaler.step(optimizer)
|
||||
scaler.update()
|
||||
optimizer.zero_grad(set_to_none=True)
|
||||
torch.mps.synchronize()
|
||||
empty_device_cache('mps')
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
unittest.main()
|
||||
|
|
@ -16,7 +16,6 @@ import warnings
|
|||
import torch
|
||||
import torch.nn.functional as F
|
||||
import torch.distributed as dist
|
||||
from contextlib import nullcontext
|
||||
from torch import optim
|
||||
from torch.nn.parallel import DistributedDataParallel
|
||||
from torch.utils.data import DataLoader, DistributedSampler
|
||||
|
|
@ -24,7 +23,11 @@ from torch.optim.lr_scheduler import CosineAnnealingLR
|
|||
from transformers import AutoTokenizer
|
||||
from model.model_minimind import MiniMindConfig, MiniMindForCausalLM
|
||||
from dataset.lm_dataset import AgentRLDataset
|
||||
from trainer.trainer_utils import Logger, is_main_process, lm_checkpoint, init_distributed_mode, setup_seed, SkipBatchSampler, init_model, LMForRewardModel
|
||||
from trainer.trainer_utils import (Logger, is_main_process, lm_checkpoint,
|
||||
init_distributed_mode, setup_seed,
|
||||
SkipBatchSampler, init_model, LMForRewardModel,
|
||||
get_default_device, get_autocast_context,
|
||||
get_grad_scaler, should_pin_memory)
|
||||
from trainer.rollout_engine import create_rollout_engine, compute_per_token_logps
|
||||
|
||||
warnings.filterwarnings('ignore')
|
||||
|
|
@ -336,11 +339,13 @@ def rl_train_epoch(epoch, loader, iters, rollout_engine, ref_model, reward_model
|
|||
policy_loss = (((per_token_loss * completion_mask).sum(dim=1)[valid_rows] / token_counts[valid_rows].clamp(min=1)).mean()
|
||||
if valid_rows.any() else per_token_loss.sum() * 0.0)
|
||||
loss = (policy_loss + aux_loss) / args.accumulation_steps
|
||||
loss.backward()
|
||||
scaler.scale(loss).backward()
|
||||
|
||||
if step % args.accumulation_steps == 0:
|
||||
if args.grad_clip > 0: torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip)
|
||||
optimizer.step(); scheduler.step(); optimizer.zero_grad()
|
||||
if args.grad_clip > 0:
|
||||
scaler.unscale_(optimizer)
|
||||
torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip)
|
||||
scaler.step(optimizer); scaler.update(); scheduler.step(); optimizer.zero_grad()
|
||||
|
||||
if step % args.log_interval == 0 or step == iters:
|
||||
pl = loss.item() * args.accumulation_steps
|
||||
|
|
@ -363,7 +368,8 @@ def rl_train_epoch(epoch, loader, iters, rollout_engine, ref_model, reward_model
|
|||
state_dict = raw_model.state_dict()
|
||||
torch.save({k: v.half().cpu() for k, v in state_dict.items()}, ckp)
|
||||
lm_checkpoint(lm_config, weight=args.save_weight, model=model, optimizer=optimizer,
|
||||
epoch=epoch, step=step, wandb=wandb, save_dir='../checkpoints', scheduler=scheduler)
|
||||
epoch=epoch, step=step, wandb=wandb, save_dir='../checkpoints',
|
||||
scheduler=scheduler, scaler=scaler)
|
||||
model.train()
|
||||
del state_dict
|
||||
|
||||
|
|
@ -373,8 +379,10 @@ def rl_train_epoch(epoch, loader, iters, rollout_engine, ref_model, reward_model
|
|||
del completions, rewards, grouped_rewards, mean_r, std_r, advantages, completion_mask
|
||||
|
||||
if last_step > start_step and last_step % args.accumulation_steps != 0:
|
||||
if args.grad_clip > 0: torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip)
|
||||
optimizer.step(); scheduler.step(); optimizer.zero_grad()
|
||||
if args.grad_clip > 0:
|
||||
scaler.unscale_(optimizer)
|
||||
torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip)
|
||||
scaler.step(optimizer); scaler.update(); scheduler.step(); optimizer.zero_grad()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
|
|
@ -384,8 +392,8 @@ if __name__ == "__main__":
|
|||
parser.add_argument("--epochs", type=int, default=1, help="训练轮数")
|
||||
parser.add_argument("--batch_size", type=int, default=2, help="批次大小")
|
||||
parser.add_argument("--learning_rate", type=float, default=3e-7, help="学习率")
|
||||
parser.add_argument("--device", type=str, default="cuda:0" if torch.cuda.is_available() else "cpu", help="训练设备")
|
||||
parser.add_argument("--dtype", type=str, default="bfloat16", help="数据类型 bfloat16/float16")
|
||||
parser.add_argument("--device", type=str, default=get_default_device(), help="训练设备 (cuda:0/mps/cpu)")
|
||||
parser.add_argument("--dtype", type=str, default="bfloat16", choices=["float32", "float16", "bfloat16"], help="训练数据类型")
|
||||
parser.add_argument("--num_workers", type=int, default=8, help="数据加载线程数")
|
||||
parser.add_argument("--accumulation_steps", type=int, default=1, help="梯度累积步数")
|
||||
parser.add_argument("--grad_clip", type=float, default=1.0, help="梯度裁剪阈值")
|
||||
|
|
@ -418,7 +426,7 @@ if __name__ == "__main__":
|
|||
parser.add_argument("--sglang_shared_path", type=str, default="./sglang_ckpt_agent", help="SGLang共享存储路径")
|
||||
args = parser.parse_args()
|
||||
|
||||
local_rank = init_distributed_mode()
|
||||
local_rank = init_distributed_mode(args.device)
|
||||
if dist.is_initialized(): args.device = f"cuda:{local_rank}"
|
||||
setup_seed(42 + (dist.get_rank() if dist.is_initialized() else 0))
|
||||
|
||||
|
|
@ -427,9 +435,8 @@ if __name__ == "__main__":
|
|||
max_seq_len=args.max_seq_len + args.max_gen_len, use_moe=bool(args.use_moe))
|
||||
ckp_data = lm_checkpoint(lm_config, weight=args.save_weight, save_dir='../checkpoints') if args.from_resume == 1 else None
|
||||
|
||||
device_type = "cuda" if "cuda" in args.device else "cpu"
|
||||
dtype = torch.bfloat16 if args.dtype == "bfloat16" else torch.float16
|
||||
autocast_ctx = nullcontext() if device_type == "cpu" else torch.cuda.amp.autocast(dtype=dtype)
|
||||
autocast_ctx = get_autocast_context(args.device, args.dtype)
|
||||
scaler = get_grad_scaler(args.device, args.dtype)
|
||||
|
||||
wandb = None
|
||||
if args.use_wandb and is_main_process():
|
||||
|
|
@ -470,6 +477,7 @@ if __name__ == "__main__":
|
|||
model.load_state_dict(ckp_data['model'])
|
||||
optimizer.load_state_dict(ckp_data['optimizer'])
|
||||
scheduler.load_state_dict(ckp_data['scheduler'])
|
||||
if 'scaler' in ckp_data: scaler.load_state_dict(ckp_data['scaler'])
|
||||
start_epoch = ckp_data['epoch']
|
||||
start_step = ckp_data.get('step', 0)
|
||||
|
||||
|
|
@ -486,7 +494,8 @@ if __name__ == "__main__":
|
|||
setup_seed(42 + epoch); indices = torch.randperm(len(train_ds)).tolist()
|
||||
skip = start_step if (epoch == start_epoch and start_step > 0) else 0
|
||||
batch_sampler = SkipBatchSampler(train_sampler or indices, args.batch_size, skip)
|
||||
loader = DataLoader(train_ds, batch_sampler=batch_sampler, num_workers=args.num_workers, pin_memory=True, collate_fn=collate_fn)
|
||||
loader = DataLoader(train_ds, batch_sampler=batch_sampler, num_workers=args.num_workers,
|
||||
pin_memory=should_pin_memory(args.device), collate_fn=collate_fn)
|
||||
if skip > 0:
|
||||
Logger(f'Epoch [{epoch+1}/{args.epochs}]: skip {start_step} steps')
|
||||
rl_train_epoch(epoch, loader, len(loader) + skip, rollout_engine, ref_model, reward_model, start_step, wandb, use_sglang = (args.rollout_engine == "sglang"))
|
||||
|
|
|
|||
|
|
@ -11,13 +11,16 @@ import warnings
|
|||
import torch
|
||||
import torch.nn.functional as F
|
||||
import torch.distributed as dist
|
||||
from contextlib import nullcontext
|
||||
from torch import optim
|
||||
from torch.nn.parallel import DistributedDataParallel
|
||||
from torch.utils.data import DataLoader, DistributedSampler
|
||||
from model.model_minimind import MiniMindConfig
|
||||
from dataset.lm_dataset import SFTDataset
|
||||
from trainer.trainer_utils import get_lr, Logger, is_main_process, lm_checkpoint, init_distributed_mode, setup_seed, init_model, SkipBatchSampler
|
||||
from trainer.trainer_utils import (get_lr, Logger, is_main_process, lm_checkpoint,
|
||||
init_distributed_mode, setup_seed, init_model,
|
||||
SkipBatchSampler, get_default_device,
|
||||
get_autocast_context, get_grad_scaler,
|
||||
should_pin_memory)
|
||||
|
||||
warnings.filterwarnings('ignore')
|
||||
|
||||
|
|
@ -151,8 +154,8 @@ if __name__ == "__main__":
|
|||
parser.add_argument("--epochs", type=int, default=6, help="训练轮数")
|
||||
parser.add_argument("--batch_size", type=int, default=32, help="batch size")
|
||||
parser.add_argument("--learning_rate", type=float, default=5e-6, help="初始学习率")
|
||||
parser.add_argument("--device", type=str, default="cuda:0" if torch.cuda.is_available() else "cpu", help="训练设备")
|
||||
parser.add_argument("--dtype", type=str, default="bfloat16", help="混合精度类型")
|
||||
parser.add_argument("--device", type=str, default=get_default_device(), help="训练设备 (cuda:0/mps/cpu)")
|
||||
parser.add_argument("--dtype", type=str, default="bfloat16", choices=["float32", "float16", "bfloat16"], help="混合精度类型")
|
||||
parser.add_argument("--num_workers", type=int, default=8, help="数据加载线程数")
|
||||
parser.add_argument("--accumulation_steps", type=int, default=1, help="梯度累积步数")
|
||||
parser.add_argument("--grad_clip", type=float, default=1.0, help="梯度裁剪阈值")
|
||||
|
|
@ -177,7 +180,7 @@ if __name__ == "__main__":
|
|||
args = parser.parse_args()
|
||||
|
||||
# ========== 1. 初始化环境和随机种子 ==========
|
||||
local_rank = init_distributed_mode()
|
||||
local_rank = init_distributed_mode(args.device)
|
||||
if dist.is_initialized(): args.device = f"cuda:{local_rank}"
|
||||
setup_seed(42 + (dist.get_rank() if dist.is_initialized() else 0))
|
||||
|
||||
|
|
@ -188,9 +191,8 @@ if __name__ == "__main__":
|
|||
ckp_data = lm_checkpoint(lm_config_student, weight=args.save_weight, save_dir='../checkpoints') if args.from_resume==1 else None
|
||||
|
||||
# ========== 3. 设置混合精度 ==========
|
||||
device_type = "cuda" if "cuda" in args.device else "cpu"
|
||||
dtype = torch.bfloat16 if args.dtype == "bfloat16" else torch.float16
|
||||
autocast_ctx = nullcontext() if device_type == "cpu" else torch.cuda.amp.autocast(dtype=dtype)
|
||||
autocast_ctx = get_autocast_context(args.device, args.dtype)
|
||||
scaler = get_grad_scaler(args.device, args.dtype)
|
||||
|
||||
# ========== 4. 配wandb ==========
|
||||
wandb = None
|
||||
|
|
@ -210,7 +212,6 @@ if __name__ == "__main__":
|
|||
Logger(f'教师模型总参数量：{sum(p.numel() for p in teacher_model.parameters()) / 1e6:.3f} M')
|
||||
train_ds = SFTDataset(args.data_path, tokenizer, max_length=args.max_seq_len)
|
||||
train_sampler = DistributedSampler(train_ds) if dist.is_initialized() else None
|
||||
scaler = torch.cuda.amp.GradScaler(enabled=(args.dtype == 'float16'))
|
||||
optimizer = optim.AdamW(model.parameters(), lr=args.learning_rate)
|
||||
|
||||
# ========== 6. 从ckp恢复状态 ==========
|
||||
|
|
@ -235,7 +236,8 @@ if __name__ == "__main__":
|
|||
setup_seed(42 + epoch); indices = torch.randperm(len(train_ds)).tolist()
|
||||
skip = start_step if (epoch == start_epoch and start_step > 0) else 0
|
||||
batch_sampler = SkipBatchSampler(train_sampler or indices, args.batch_size, skip)
|
||||
loader = DataLoader(train_ds, batch_sampler=batch_sampler, num_workers=args.num_workers, pin_memory=True)
|
||||
loader = DataLoader(train_ds, batch_sampler=batch_sampler, num_workers=args.num_workers,
|
||||
pin_memory=should_pin_memory(args.device))
|
||||
if skip > 0:
|
||||
Logger(f'Epoch [{epoch + 1}/{args.epochs}]: 跳过前{start_step}个step，从step {start_step + 1}开始')
|
||||
train_epoch(epoch, loader, len(loader) + skip, teacher_model, lm_config_student, start_step, wandb, args.alpha, args.temperature)
|
||||
|
|
@ -245,4 +247,4 @@ if __name__ == "__main__":
|
|||
# ========== 9. 清理分布进程 ==========
|
||||
if dist.is_initialized():
|
||||
dist.barrier()
|
||||
dist.destroy_process_group()
|
||||
dist.destroy_process_group()
|
||||
|
|
|
|||
|
|
@ -11,13 +11,16 @@ import warnings
|
|||
import torch
|
||||
import torch.nn.functional as F
|
||||
import torch.distributed as dist
|
||||
from contextlib import nullcontext
|
||||
from torch import optim
|
||||
from torch.nn.parallel import DistributedDataParallel
|
||||
from torch.utils.data import DataLoader, DistributedSampler
|
||||
from model.model_minimind import MiniMindConfig
|
||||
from dataset.lm_dataset import DPODataset
|
||||
from trainer.trainer_utils import get_lr, Logger, is_main_process, lm_checkpoint, init_distributed_mode, setup_seed, init_model, SkipBatchSampler
|
||||
from trainer.trainer_utils import (get_lr, Logger, is_main_process, lm_checkpoint,
|
||||
init_distributed_mode, setup_seed, init_model,
|
||||
SkipBatchSampler, get_default_device,
|
||||
get_autocast_context, get_grad_scaler,
|
||||
should_pin_memory)
|
||||
|
||||
warnings.filterwarnings('ignore')
|
||||
|
||||
|
|
@ -135,8 +138,8 @@ if __name__ == "__main__":
|
|||
parser.add_argument("--epochs", type=int, default=1, help="训练轮数")
|
||||
parser.add_argument("--batch_size", type=int, default=4, help="batch size")
|
||||
parser.add_argument("--learning_rate", type=float, default=4e-8, help="初始学习率（建议<=5e-8避免遗忘）")
|
||||
parser.add_argument("--device", type=str, default="cuda:0" if torch.cuda.is_available() else "cpu", help="训练设备")
|
||||
parser.add_argument("--dtype", type=str, default="bfloat16", help="混合精度类型")
|
||||
parser.add_argument("--device", type=str, default=get_default_device(), help="训练设备 (cuda:0/mps/cpu)")
|
||||
parser.add_argument("--dtype", type=str, default="bfloat16", choices=["float32", "float16", "bfloat16"], help="混合精度类型")
|
||||
parser.add_argument("--num_workers", type=int, default=8, help="数据加载线程数")
|
||||
parser.add_argument("--accumulation_steps", type=int, default=1, help="梯度累积步数")
|
||||
parser.add_argument("--grad_clip", type=float, default=1.0, help="梯度裁剪阈值")
|
||||
|
|
@ -156,7 +159,7 @@ if __name__ == "__main__":
|
|||
args = parser.parse_args()
|
||||
|
||||
# ========== 1. 初始化环境和随机种子 ==========
|
||||
local_rank = init_distributed_mode()
|
||||
local_rank = init_distributed_mode(args.device)
|
||||
if dist.is_initialized(): args.device = f"cuda:{local_rank}"
|
||||
setup_seed(42 + (dist.get_rank() if dist.is_initialized() else 0))
|
||||
|
||||
|
|
@ -166,9 +169,8 @@ if __name__ == "__main__":
|
|||
ckp_data = lm_checkpoint(lm_config, weight=args.save_weight, save_dir='../checkpoints') if args.from_resume==1 else None
|
||||
|
||||
# ========== 3. 设置混合精度 ==========
|
||||
device_type = "cuda" if "cuda" in args.device else "cpu"
|
||||
dtype = torch.bfloat16 if args.dtype == "bfloat16" else torch.float16
|
||||
autocast_ctx = nullcontext() if device_type == "cpu" else torch.cuda.amp.autocast(dtype=dtype)
|
||||
autocast_ctx = get_autocast_context(args.device, args.dtype)
|
||||
scaler = get_grad_scaler(args.device, args.dtype)
|
||||
|
||||
# ========== 4. 配wandb ==========
|
||||
wandb = None
|
||||
|
|
@ -190,7 +192,6 @@ if __name__ == "__main__":
|
|||
|
||||
train_ds = DPODataset(args.data_path, tokenizer, max_length=args.max_seq_len)
|
||||
train_sampler = DistributedSampler(train_ds) if dist.is_initialized() else None
|
||||
scaler = torch.cuda.amp.GradScaler(enabled=(args.dtype == 'float16'))
|
||||
optimizer = optim.AdamW(model.parameters(), lr=args.learning_rate)
|
||||
|
||||
# ========== 6. 从ckp恢复状态 ==========
|
||||
|
|
@ -215,7 +216,8 @@ if __name__ == "__main__":
|
|||
setup_seed(42 + epoch); indices = torch.randperm(len(train_ds)).tolist()
|
||||
skip = start_step if (epoch == start_epoch and start_step > 0) else 0
|
||||
batch_sampler = SkipBatchSampler(train_sampler or indices, args.batch_size, skip)
|
||||
loader = DataLoader(train_ds, batch_sampler=batch_sampler, num_workers=args.num_workers, pin_memory=True)
|
||||
loader = DataLoader(train_ds, batch_sampler=batch_sampler, num_workers=args.num_workers,
|
||||
pin_memory=should_pin_memory(args.device))
|
||||
if skip > 0:
|
||||
Logger(f'Epoch [{epoch + 1}/{args.epochs}]: 跳过前{start_step}个step，从step {start_step + 1}开始')
|
||||
train_epoch(epoch, loader, len(loader) + skip, ref_model, lm_config, start_step, wandb, args.beta)
|
||||
|
|
@ -225,4 +227,4 @@ if __name__ == "__main__":
|
|||
# ========== 9. 清理分布进程 ==========
|
||||
if dist.is_initialized():
|
||||
dist.barrier()
|
||||
dist.destroy_process_group()
|
||||
dist.destroy_process_group()
|
||||
|
|
|
|||
|
|
@ -10,13 +10,16 @@ import time
|
|||
import warnings
|
||||
import torch
|
||||
import torch.distributed as dist
|
||||
from contextlib import nullcontext
|
||||
from torch import optim, nn
|
||||
from torch.nn.parallel import DistributedDataParallel
|
||||
from torch.utils.data import DataLoader, DistributedSampler
|
||||
from model.model_minimind import MiniMindConfig
|
||||
from dataset.lm_dataset import SFTDataset
|
||||
from trainer.trainer_utils import get_lr, Logger, is_main_process, lm_checkpoint, init_distributed_mode, setup_seed, init_model, SkipBatchSampler
|
||||
from trainer.trainer_utils import (get_lr, Logger, is_main_process, lm_checkpoint,
|
||||
init_distributed_mode, setup_seed, init_model,
|
||||
SkipBatchSampler, get_default_device,
|
||||
get_autocast_context, get_grad_scaler,
|
||||
should_pin_memory)
|
||||
|
||||
warnings.filterwarnings('ignore')
|
||||
|
||||
|
|
@ -88,8 +91,8 @@ if __name__ == "__main__":
|
|||
parser.add_argument("--epochs", type=int, default=2, help="训练轮数")
|
||||
parser.add_argument("--batch_size", type=int, default=16, help="batch size")
|
||||
parser.add_argument("--learning_rate", type=float, default=1e-5, help="初始学习率")
|
||||
parser.add_argument("--device", type=str, default="cuda:0" if torch.cuda.is_available() else "cpu", help="训练设备")
|
||||
parser.add_argument("--dtype", type=str, default="bfloat16", help="混合精度类型")
|
||||
parser.add_argument("--device", type=str, default=get_default_device(), help="训练设备 (cuda:0/mps/cpu)")
|
||||
parser.add_argument("--dtype", type=str, default="bfloat16", choices=["float32", "float16", "bfloat16"], help="混合精度类型")
|
||||
parser.add_argument("--num_workers", type=int, default=8, help="数据加载线程数")
|
||||
parser.add_argument("--accumulation_steps", type=int, default=1, help="梯度累积步数")
|
||||
parser.add_argument("--grad_clip", type=float, default=1.0, help="梯度裁剪阈值")
|
||||
|
|
@ -109,7 +112,7 @@ if __name__ == "__main__":
|
|||
args = parser.parse_args()
|
||||
|
||||
# ========== 1. 初始化环境和随机种子 ==========
|
||||
local_rank = init_distributed_mode()
|
||||
local_rank = init_distributed_mode(args.device)
|
||||
if dist.is_initialized(): args.device = f"cuda:{local_rank}"
|
||||
setup_seed(args.seed + (dist.get_rank() if dist.is_initialized() else 0))
|
||||
|
||||
|
|
@ -119,9 +122,8 @@ if __name__ == "__main__":
|
|||
ckp_data = lm_checkpoint(lm_config, weight=args.save_weight, save_dir='../checkpoints') if args.from_resume==1 else None
|
||||
|
||||
# ========== 3. 设置混合精度 ==========
|
||||
device_type = "cuda" if "cuda" in args.device else "cpu"
|
||||
dtype = torch.bfloat16 if args.dtype == "bfloat16" else torch.float16
|
||||
autocast_ctx = nullcontext() if device_type == "cpu" else torch.cuda.amp.autocast(dtype=dtype)
|
||||
autocast_ctx = get_autocast_context(args.device, args.dtype)
|
||||
scaler = get_grad_scaler(args.device, args.dtype)
|
||||
|
||||
# ========== 4. 配wandb ==========
|
||||
wandb = None
|
||||
|
|
@ -136,7 +138,6 @@ if __name__ == "__main__":
|
|||
model, tokenizer = init_model(lm_config, args.from_weight, device=args.device)
|
||||
train_ds = SFTDataset(args.data_path, tokenizer, max_length=args.max_seq_len)
|
||||
train_sampler = DistributedSampler(train_ds) if dist.is_initialized() else None
|
||||
scaler = torch.cuda.amp.GradScaler(enabled=(args.dtype == 'float16'))
|
||||
optimizer = optim.AdamW(model.parameters(), lr=args.learning_rate)
|
||||
|
||||
# ========== 6. 从ckp恢复状态 ==========
|
||||
|
|
@ -161,7 +162,8 @@ if __name__ == "__main__":
|
|||
setup_seed(args.seed + epoch); indices = torch.randperm(len(train_ds)).tolist()
|
||||
skip = start_step if (epoch == start_epoch and start_step > 0) else 0
|
||||
batch_sampler = SkipBatchSampler(train_sampler or indices, args.batch_size, skip)
|
||||
loader = DataLoader(train_ds, batch_sampler=batch_sampler, num_workers=args.num_workers, pin_memory=True)
|
||||
loader = DataLoader(train_ds, batch_sampler=batch_sampler, num_workers=args.num_workers,
|
||||
pin_memory=should_pin_memory(args.device))
|
||||
if skip > 0:
|
||||
Logger(f'Epoch [{epoch + 1}/{args.epochs}]: 跳过前{start_step}个step，从step {start_step + 1}开始')
|
||||
train_epoch(epoch, loader, len(loader) + skip, start_step, wandb)
|
||||
|
|
@ -171,4 +173,4 @@ if __name__ == "__main__":
|
|||
# ========== 9. 清理分布进程 ==========
|
||||
if dist.is_initialized():
|
||||
dist.barrier()
|
||||
dist.destroy_process_group()
|
||||
dist.destroy_process_group()
|
||||
|
|
|
|||
|
|
@ -14,7 +14,6 @@ import torch
|
|||
import torch.nn.functional as F
|
||||
import torch.distributed as dist
|
||||
from transformers import AutoTokenizer
|
||||
from contextlib import nullcontext
|
||||
from torch import optim
|
||||
from torch.nn.parallel import DistributedDataParallel
|
||||
from torch.utils.data import DataLoader, DistributedSampler
|
||||
|
|
@ -22,7 +21,11 @@ from torch.optim.lr_scheduler import CosineAnnealingLR
|
|||
from transformers import AutoModel
|
||||
from model.model_minimind import MiniMindConfig, MiniMindForCausalLM
|
||||
from dataset.lm_dataset import RLAIFDataset
|
||||
from trainer.trainer_utils import Logger, is_main_process, lm_checkpoint, init_distributed_mode, setup_seed, SkipBatchSampler, init_model, LMForRewardModel
|
||||
from trainer.trainer_utils import (Logger, is_main_process, lm_checkpoint,
|
||||
init_distributed_mode, setup_seed,
|
||||
SkipBatchSampler, init_model, LMForRewardModel,
|
||||
get_default_device, get_autocast_context,
|
||||
get_grad_scaler, should_pin_memory)
|
||||
from trainer.rollout_engine import create_rollout_engine
|
||||
|
||||
warnings.filterwarnings('ignore')
|
||||
|
|
@ -143,12 +146,14 @@ def grpo_train_epoch(epoch, loader, iters, rollout_engine, ref_model, reward_mod
|
|||
per_token_loss = -(torch.min(per_token_loss1, per_token_loss2) - args.beta * per_token_kl)
|
||||
policy_loss = ((per_token_loss * completion_mask).sum(dim=1) / completion_mask.sum(dim=1).clamp(min=1)).mean()
|
||||
loss = (policy_loss + aux_loss) / args.accumulation_steps # scalar
|
||||
loss.backward()
|
||||
scaler.scale(loss).backward()
|
||||
|
||||
if step % args.accumulation_steps == 0:
|
||||
if args.grad_clip > 0:
|
||||
scaler.unscale_(optimizer)
|
||||
torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip)
|
||||
optimizer.step()
|
||||
scaler.step(optimizer)
|
||||
scaler.update()
|
||||
scheduler.step()
|
||||
optimizer.zero_grad()
|
||||
|
||||
|
|
@ -187,7 +192,8 @@ def grpo_train_epoch(epoch, loader, iters, rollout_engine, ref_model, reward_mod
|
|||
state_dict = raw_model.state_dict()
|
||||
torch.save({k: v.half().cpu() for k, v in state_dict.items()}, ckp)
|
||||
lm_checkpoint(lm_config, weight=args.save_weight, model=model, optimizer=optimizer,
|
||||
epoch=epoch, step=step, wandb=wandb, save_dir='../checkpoints', scheduler=scheduler)
|
||||
epoch=epoch, step=step, wandb=wandb, save_dir='../checkpoints',
|
||||
scheduler=scheduler, scaler=scaler)
|
||||
model.train()
|
||||
del state_dict
|
||||
|
||||
|
|
@ -198,8 +204,10 @@ def grpo_train_epoch(epoch, loader, iters, rollout_engine, ref_model, reward_mod
|
|||
|
||||
if step > start_step and step % args.accumulation_steps != 0:
|
||||
if args.grad_clip > 0:
|
||||
scaler.unscale_(optimizer)
|
||||
torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip)
|
||||
optimizer.step()
|
||||
scaler.step(optimizer)
|
||||
scaler.update()
|
||||
scheduler.step()
|
||||
optimizer.zero_grad()
|
||||
|
||||
|
|
@ -211,8 +219,8 @@ if __name__ == "__main__":
|
|||
parser.add_argument("--epochs", type=int, default=1, help="训练轮数")
|
||||
parser.add_argument("--batch_size", type=int, default=2, help="batch size")
|
||||
parser.add_argument("--learning_rate", type=float, default=3e-7, help="初始学习率")
|
||||
parser.add_argument("--device", type=str, default="cuda:0" if torch.cuda.is_available() else "cpu", help="训练设备")
|
||||
parser.add_argument("--dtype", type=str, default="bfloat16", help="混合精度类型")
|
||||
parser.add_argument("--device", type=str, default=get_default_device(), help="训练设备 (cuda:0/mps/cpu)")
|
||||
parser.add_argument("--dtype", type=str, default="bfloat16", choices=["float32", "float16", "bfloat16"], help="混合精度类型")
|
||||
parser.add_argument("--num_workers", type=int, default=8, help="数据加载线程数")
|
||||
parser.add_argument("--accumulation_steps", type=int, default=1, help="梯度累积步数")
|
||||
parser.add_argument("--grad_clip", type=float, default=1.0, help="梯度裁剪阈值")
|
||||
|
|
@ -245,7 +253,7 @@ if __name__ == "__main__":
|
|||
args = parser.parse_args()
|
||||
|
||||
# ========== 1. 初始化环境和随机种子 ==========
|
||||
local_rank = init_distributed_mode()
|
||||
local_rank = init_distributed_mode(args.device)
|
||||
if dist.is_initialized(): args.device = f"cuda:{local_rank}"
|
||||
setup_seed(42 + (dist.get_rank() if dist.is_initialized() else 0))
|
||||
|
||||
|
|
@ -256,9 +264,8 @@ if __name__ == "__main__":
|
|||
ckp_data = lm_checkpoint(lm_config, weight=args.save_weight, save_dir='../checkpoints') if args.from_resume==1 else None
|
||||
|
||||
# ========== 3. 设置混合精度 ==========
|
||||
device_type = "cuda" if "cuda" in args.device else "cpu"
|
||||
dtype = torch.bfloat16 if args.dtype == "bfloat16" else torch.float16
|
||||
autocast_ctx = nullcontext() if device_type == "cpu" else torch.cuda.amp.autocast(dtype=dtype)
|
||||
autocast_ctx = get_autocast_context(args.device, args.dtype)
|
||||
scaler = get_grad_scaler(args.device, args.dtype)
|
||||
|
||||
# ========== 4. 配wandb ==========
|
||||
wandb = None
|
||||
|
|
@ -304,6 +311,7 @@ if __name__ == "__main__":
|
|||
model.load_state_dict(ckp_data['model'])
|
||||
optimizer.load_state_dict(ckp_data['optimizer'])
|
||||
scheduler.load_state_dict(ckp_data['scheduler'])
|
||||
if 'scaler' in ckp_data: scaler.load_state_dict(ckp_data['scaler'])
|
||||
start_epoch = ckp_data['epoch']
|
||||
start_step = ckp_data.get('step', 0)
|
||||
|
||||
|
|
@ -322,7 +330,8 @@ if __name__ == "__main__":
|
|||
setup_seed(42 + epoch); indices = torch.randperm(len(train_ds)).tolist()
|
||||
skip = start_step if (epoch == start_epoch and start_step > 0) else 0
|
||||
batch_sampler = SkipBatchSampler(train_sampler or indices, args.batch_size, skip)
|
||||
loader = DataLoader(train_ds, batch_sampler=batch_sampler, num_workers=args.num_workers, pin_memory=True)
|
||||
loader = DataLoader(train_ds, batch_sampler=batch_sampler, num_workers=args.num_workers,
|
||||
pin_memory=should_pin_memory(args.device))
|
||||
if skip > 0:
|
||||
Logger(f'Epoch [{epoch + 1}/{args.epochs}]: 跳过前{start_step}个step，从step {start_step + 1}开始')
|
||||
grpo_train_epoch(epoch, loader, len(loader) + skip, rollout_engine, ref_model, reward_model, start_step, wandb, use_sglang = (args.rollout_engine == "sglang"))
|
||||
|
|
@ -332,4 +341,4 @@ if __name__ == "__main__":
|
|||
# ========== 9. 清理分布进程 ==========
|
||||
if dist.is_initialized():
|
||||
dist.barrier()
|
||||
dist.destroy_process_group()
|
||||
dist.destroy_process_group()
|
||||
|
|
|
|||
|
|
@ -10,14 +10,17 @@ import time
|
|||
import warnings
|
||||
import torch
|
||||
import torch.distributed as dist
|
||||
from contextlib import nullcontext
|
||||
from torch import optim, nn
|
||||
from torch.nn.parallel import DistributedDataParallel
|
||||
from torch.utils.data import DataLoader, DistributedSampler
|
||||
from model.model_minimind import MiniMindConfig
|
||||
from dataset.lm_dataset import SFTDataset
|
||||
from model.model_lora import save_lora, apply_lora
|
||||
from trainer.trainer_utils import get_lr, Logger, is_main_process, lm_checkpoint, init_distributed_mode, setup_seed, init_model, SkipBatchSampler
|
||||
from trainer.trainer_utils import (get_lr, Logger, is_main_process, lm_checkpoint,
|
||||
init_distributed_mode, setup_seed, init_model,
|
||||
SkipBatchSampler, get_default_device,
|
||||
get_autocast_context, get_grad_scaler,
|
||||
should_pin_memory)
|
||||
|
||||
warnings.filterwarnings('ignore')
|
||||
|
||||
|
|
@ -82,8 +85,8 @@ if __name__ == "__main__":
|
|||
parser.add_argument("--epochs", type=int, default=10, help="训练轮数")
|
||||
parser.add_argument("--batch_size", type=int, default=32, help="batch size")
|
||||
parser.add_argument("--learning_rate", type=float, default=1e-4, help="初始学习率")
|
||||
parser.add_argument("--device", type=str, default="cuda:0" if torch.cuda.is_available() else "cpu", help="训练设备")
|
||||
parser.add_argument("--dtype", type=str, default="bfloat16", help="混合精度类型")
|
||||
parser.add_argument("--device", type=str, default=get_default_device(), help="训练设备 (cuda:0/mps/cpu)")
|
||||
parser.add_argument("--dtype", type=str, default="bfloat16", choices=["float32", "float16", "bfloat16"], help="混合精度类型")
|
||||
parser.add_argument("--num_workers", type=int, default=8, help="数据加载线程数")
|
||||
parser.add_argument("--accumulation_steps", type=int, default=1, help="梯度累积步数")
|
||||
parser.add_argument("--grad_clip", type=float, default=1.0, help="梯度裁剪阈值")
|
||||
|
|
@ -102,7 +105,7 @@ if __name__ == "__main__":
|
|||
args = parser.parse_args()
|
||||
|
||||
# ========== 1. 初始化环境和随机种子 ==========
|
||||
local_rank = init_distributed_mode()
|
||||
local_rank = init_distributed_mode(args.device)
|
||||
if dist.is_initialized(): args.device = f"cuda:{local_rank}"
|
||||
setup_seed(42 + (dist.get_rank() if dist.is_initialized() else 0))
|
||||
|
||||
|
|
@ -112,9 +115,8 @@ if __name__ == "__main__":
|
|||
ckp_data = lm_checkpoint(lm_config, weight=args.lora_name, save_dir='../checkpoints') if args.from_resume==1 else None
|
||||
|
||||
# ========== 3. 设置混合精度 ==========
|
||||
device_type = "cuda" if "cuda" in args.device else "cpu"
|
||||
dtype = torch.bfloat16 if args.dtype == "bfloat16" else torch.float16
|
||||
autocast_ctx = nullcontext() if device_type == "cpu" else torch.cuda.amp.autocast(dtype=dtype)
|
||||
autocast_ctx = get_autocast_context(args.device, args.dtype)
|
||||
scaler = get_grad_scaler(args.device, args.dtype)
|
||||
|
||||
# ========== 4. 配wandb ==========
|
||||
wandb = None
|
||||
|
|
@ -148,7 +150,6 @@ if __name__ == "__main__":
|
|||
# ========== 6. 定义数据和优化器 ==========
|
||||
train_ds = SFTDataset(args.data_path, tokenizer, max_length=args.max_seq_len)
|
||||
train_sampler = DistributedSampler(train_ds) if dist.is_initialized() else None
|
||||
scaler = torch.cuda.amp.GradScaler(enabled=(args.dtype == 'float16'))
|
||||
optimizer = optim.AdamW(lora_params, lr=args.learning_rate)
|
||||
|
||||
# ========== 7. 从ckp恢复状态 ==========
|
||||
|
|
@ -173,7 +174,8 @@ if __name__ == "__main__":
|
|||
setup_seed(42 + epoch); indices = torch.randperm(len(train_ds)).tolist()
|
||||
skip = start_step if (epoch == start_epoch and start_step > 0) else 0
|
||||
batch_sampler = SkipBatchSampler(train_sampler or indices, args.batch_size, skip)
|
||||
loader = DataLoader(train_ds, batch_sampler=batch_sampler, num_workers=args.num_workers, pin_memory=True)
|
||||
loader = DataLoader(train_ds, batch_sampler=batch_sampler, num_workers=args.num_workers,
|
||||
pin_memory=should_pin_memory(args.device))
|
||||
if skip > 0:
|
||||
Logger(f'Epoch [{epoch + 1}/{args.epochs}]: 跳过前{start_step}个step，从step {start_step + 1}开始')
|
||||
train_epoch(epoch, loader, len(loader) + skip, lora_params, start_step, wandb)
|
||||
|
|
@ -183,4 +185,4 @@ if __name__ == "__main__":
|
|||
# ========== 10. 清理分布进程 ==========
|
||||
if dist.is_initialized():
|
||||
dist.barrier()
|
||||
dist.destroy_process_group()
|
||||
dist.destroy_process_group()
|
||||
|
|
|
|||
|
|
@ -12,7 +12,6 @@ import warnings
|
|||
import torch
|
||||
import torch.distributed as dist
|
||||
import torch.nn.functional as F
|
||||
from contextlib import nullcontext
|
||||
from torch import optim, nn
|
||||
from torch.nn.parallel import DistributedDataParallel
|
||||
from torch.utils.data import DataLoader, DistributedSampler
|
||||
|
|
@ -20,7 +19,11 @@ from torch.nn.utils import clip_grad_norm_
|
|||
from torch.optim.lr_scheduler import CosineAnnealingLR
|
||||
from model.model_minimind import MiniMindConfig, MiniMindForCausalLM
|
||||
from dataset.lm_dataset import RLAIFDataset
|
||||
from trainer.trainer_utils import Logger, is_main_process, lm_checkpoint, init_distributed_mode, setup_seed, SkipBatchSampler, init_model, LMForRewardModel
|
||||
from trainer.trainer_utils import (Logger, is_main_process, lm_checkpoint,
|
||||
init_distributed_mode, setup_seed,
|
||||
SkipBatchSampler, init_model, LMForRewardModel,
|
||||
get_default_device, get_autocast_context,
|
||||
get_grad_scaler, should_pin_memory)
|
||||
from trainer.rollout_engine import create_rollout_engine
|
||||
|
||||
warnings.filterwarnings('ignore')
|
||||
|
|
@ -225,7 +228,7 @@ def ppo_train_epoch(epoch, loader, iters, rollout_engine, ref_model, actor_sched
|
|||
else:
|
||||
loss = (policy_loss + args.vf_coef * value_loss + aux_loss) / args.accumulation_steps
|
||||
|
||||
loss.backward()
|
||||
scaler.scale(loss).backward()
|
||||
|
||||
policy_loss_sum += policy_loss.item()
|
||||
value_loss_sum += value_loss.item()
|
||||
|
|
@ -238,20 +241,26 @@ def ppo_train_epoch(epoch, loader, iters, rollout_engine, ref_model, actor_sched
|
|||
grad_accum_step += 1
|
||||
|
||||
if grad_accum_step % args.accumulation_steps == 0:
|
||||
scaler.unscale_(actor_optimizer)
|
||||
scaler.unscale_(critic_optimizer)
|
||||
clip_grad_norm_(actor_model.parameters(), args.grad_clip)
|
||||
clip_grad_norm_(critic_model.parameters(), args.grad_clip)
|
||||
actor_optimizer.step()
|
||||
critic_optimizer.step()
|
||||
scaler.step(actor_optimizer)
|
||||
scaler.step(critic_optimizer)
|
||||
scaler.update()
|
||||
actor_scheduler.step()
|
||||
critic_scheduler.step()
|
||||
actor_optimizer.zero_grad()
|
||||
critic_optimizer.zero_grad()
|
||||
|
||||
if grad_accum_step % args.accumulation_steps != 0:
|
||||
scaler.unscale_(actor_optimizer)
|
||||
scaler.unscale_(critic_optimizer)
|
||||
clip_grad_norm_(actor_model.parameters(), args.grad_clip)
|
||||
clip_grad_norm_(critic_model.parameters(), args.grad_clip)
|
||||
actor_optimizer.step()
|
||||
critic_optimizer.step()
|
||||
scaler.step(actor_optimizer)
|
||||
scaler.step(critic_optimizer)
|
||||
scaler.update()
|
||||
actor_scheduler.step()
|
||||
critic_scheduler.step()
|
||||
actor_optimizer.zero_grad()
|
||||
|
|
@ -298,7 +307,8 @@ def ppo_train_epoch(epoch, loader, iters, rollout_engine, ref_model, actor_sched
|
|||
lm_checkpoint(lm_config, weight=args.save_weight, model=actor_model, optimizer=actor_optimizer,
|
||||
epoch=epoch, step=step, wandb=wandb, save_dir='../checkpoints',
|
||||
scheduler=actor_scheduler, critic_model=critic_model,
|
||||
critic_optimizer=critic_optimizer, critic_scheduler=critic_scheduler)
|
||||
critic_optimizer=critic_optimizer, critic_scheduler=critic_scheduler,
|
||||
scaler=scaler)
|
||||
actor_model.train()
|
||||
del actor_state
|
||||
|
||||
|
|
@ -315,8 +325,8 @@ if __name__ == "__main__":
|
|||
parser.add_argument("--batch_size", type=int, default=2, help="batch size")
|
||||
parser.add_argument("--learning_rate", type=float, default=3e-7, help="Actor学习率")
|
||||
parser.add_argument("--critic_learning_rate", type=float, default=5e-7, help="Critic学习率")
|
||||
parser.add_argument("--device", type=str, default="cuda:0" if torch.cuda.is_available() else "cpu", help="训练设备")
|
||||
parser.add_argument("--dtype", type=str, default="bfloat16", help="混合精度类型")
|
||||
parser.add_argument("--device", type=str, default=get_default_device(), help="训练设备 (cuda:0/mps/cpu)")
|
||||
parser.add_argument("--dtype", type=str, default="bfloat16", choices=["float32", "float16", "bfloat16"], help="混合精度类型")
|
||||
parser.add_argument("--num_workers", type=int, default=8, help="数据加载线程数")
|
||||
parser.add_argument("--accumulation_steps", type=int, default=1, help="梯度累积步数")
|
||||
parser.add_argument("--grad_clip", type=float, default=1.0, help="梯度裁剪阈值")
|
||||
|
|
@ -354,7 +364,7 @@ if __name__ == "__main__":
|
|||
args = parser.parse_args()
|
||||
|
||||
# ========== 1. 初始化环境和随机种子 ==========
|
||||
local_rank = init_distributed_mode()
|
||||
local_rank = init_distributed_mode(args.device)
|
||||
if dist.is_initialized(): args.device = f"cuda:{local_rank}"
|
||||
setup_seed(42 + (dist.get_rank() if dist.is_initialized() else 0))
|
||||
|
||||
|
|
@ -364,9 +374,8 @@ if __name__ == "__main__":
|
|||
ckp_data = lm_checkpoint(lm_config, weight=args.save_weight, save_dir='../checkpoints') if args.from_resume==1 else None
|
||||
|
||||
# ========== 3. 设置混合精度 ==========
|
||||
device_type = "cuda" if "cuda" in args.device else "cpu"
|
||||
dtype = torch.bfloat16 if args.dtype == "bfloat16" else torch.float16
|
||||
autocast_ctx = nullcontext() if device_type == "cpu" else torch.cuda.amp.autocast(dtype=dtype)
|
||||
autocast_ctx = get_autocast_context(args.device, args.dtype)
|
||||
scaler = get_grad_scaler(args.device, args.dtype)
|
||||
|
||||
# ========== 4. 配wandb ==========
|
||||
wandb = None
|
||||
|
|
@ -420,6 +429,7 @@ if __name__ == "__main__":
|
|||
critic_optimizer.load_state_dict(ckp_data['critic_optimizer'])
|
||||
actor_scheduler.load_state_dict(ckp_data['scheduler'])
|
||||
critic_scheduler.load_state_dict(ckp_data['critic_scheduler'])
|
||||
if 'scaler' in ckp_data: scaler.load_state_dict(ckp_data['scaler'])
|
||||
start_epoch = ckp_data['epoch']
|
||||
start_step = ckp_data.get('step', 0)
|
||||
|
||||
|
|
@ -439,7 +449,8 @@ if __name__ == "__main__":
|
|||
setup_seed(42 + epoch); indices = torch.randperm(len(train_ds)).tolist()
|
||||
skip = start_step if (epoch == start_epoch and start_step > 0) else 0
|
||||
batch_sampler = SkipBatchSampler(train_sampler or indices, args.batch_size, skip)
|
||||
loader = DataLoader(train_ds, batch_sampler=batch_sampler, num_workers=args.num_workers, pin_memory=True)
|
||||
loader = DataLoader(train_ds, batch_sampler=batch_sampler, num_workers=args.num_workers,
|
||||
pin_memory=should_pin_memory(args.device))
|
||||
if skip > 0:
|
||||
Logger(f'Epoch [{epoch + 1}/{args.epochs}]: 跳过前{start_step}个step，从step {start_step + 1}开始')
|
||||
ppo_train_epoch(epoch, loader, len(loader) + skip, rollout_engine, ref_model, actor_scheduler, critic_scheduler, reward_model, start_step, wandb)
|
||||
|
|
@ -449,4 +460,4 @@ if __name__ == "__main__":
|
|||
# ========== 9. 清理分布进程 ==========
|
||||
if dist.is_initialized():
|
||||
dist.barrier()
|
||||
dist.destroy_process_group()
|
||||
dist.destroy_process_group()
|
||||
|
|
|
|||
|
|
@ -10,13 +10,16 @@ import time
|
|||
import warnings
|
||||
import torch
|
||||
import torch.distributed as dist
|
||||
from contextlib import nullcontext
|
||||
from torch import optim, nn
|
||||
from torch.nn.parallel import DistributedDataParallel
|
||||
from torch.utils.data import DataLoader, DistributedSampler
|
||||
from model.model_minimind import MiniMindConfig
|
||||
from dataset.lm_dataset import PretrainDataset
|
||||
from trainer.trainer_utils import get_lr, Logger, is_main_process, lm_checkpoint, init_distributed_mode, setup_seed, init_model, SkipBatchSampler
|
||||
from trainer.trainer_utils import (get_lr, Logger, is_main_process, lm_checkpoint,
|
||||
init_distributed_mode, setup_seed, init_model,
|
||||
SkipBatchSampler, get_default_device,
|
||||
get_autocast_context, get_grad_scaler,
|
||||
should_pin_memory)
|
||||
|
||||
warnings.filterwarnings('ignore')
|
||||
|
||||
|
|
@ -87,8 +90,8 @@ if __name__ == "__main__":
|
|||
parser.add_argument("--epochs", type=int, default=2, help="训练轮数")
|
||||
parser.add_argument("--batch_size", type=int, default=32, help="batch size")
|
||||
parser.add_argument("--learning_rate", type=float, default=5e-4, help="初始学习率")
|
||||
parser.add_argument("--device", type=str, default="cuda:0" if torch.cuda.is_available() else "cpu", help="训练设备")
|
||||
parser.add_argument("--dtype", type=str, default="bfloat16", help="混合精度类型")
|
||||
parser.add_argument("--device", type=str, default=get_default_device(), help="训练设备 (cuda:0/mps/cpu)")
|
||||
parser.add_argument("--dtype", type=str, default="bfloat16", choices=["float32", "float16", "bfloat16"], help="混合精度类型")
|
||||
parser.add_argument("--num_workers", type=int, default=8, help="数据加载线程数")
|
||||
parser.add_argument("--accumulation_steps", type=int, default=8, help="梯度累积步数")
|
||||
parser.add_argument("--grad_clip", type=float, default=1.0, help="梯度裁剪阈值")
|
||||
|
|
@ -108,7 +111,7 @@ if __name__ == "__main__":
|
|||
args = parser.parse_args()
|
||||
|
||||
# ========== 1. 初始化环境和随机种子 ==========
|
||||
local_rank = init_distributed_mode()
|
||||
local_rank = init_distributed_mode(args.device)
|
||||
if dist.is_initialized(): args.device = f"cuda:{local_rank}"
|
||||
setup_seed(args.seed + (dist.get_rank() if dist.is_initialized() else 0))
|
||||
|
||||
|
|
@ -118,9 +121,8 @@ if __name__ == "__main__":
|
|||
ckp_data = lm_checkpoint(lm_config, weight=args.save_weight, save_dir='../checkpoints') if args.from_resume==1 else None
|
||||
|
||||
# ========== 3. 设置混合精度 ==========
|
||||
device_type = "cuda" if "cuda" in args.device else "cpu"
|
||||
dtype = torch.bfloat16 if args.dtype == "bfloat16" else torch.float16
|
||||
autocast_ctx = nullcontext() if device_type == "cpu" else torch.cuda.amp.autocast(dtype=dtype)
|
||||
autocast_ctx = get_autocast_context(args.device, args.dtype)
|
||||
scaler = get_grad_scaler(args.device, args.dtype)
|
||||
|
||||
# ========== 4. 配wandb ==========
|
||||
wandb = None
|
||||
|
|
@ -135,7 +137,6 @@ if __name__ == "__main__":
|
|||
model, tokenizer = init_model(lm_config, args.from_weight, device=args.device)
|
||||
train_ds = PretrainDataset(args.data_path, tokenizer, max_length=args.max_seq_len)
|
||||
train_sampler = DistributedSampler(train_ds) if dist.is_initialized() else None
|
||||
scaler = torch.cuda.amp.GradScaler(enabled=(args.dtype == 'float16'))
|
||||
optimizer = optim.AdamW(model.parameters(), lr=args.learning_rate)
|
||||
|
||||
# ========== 6. 从ckp恢复状态 ==========
|
||||
|
|
@ -160,7 +161,8 @@ if __name__ == "__main__":
|
|||
setup_seed(args.seed + epoch); indices = torch.randperm(len(train_ds)).tolist()
|
||||
skip = start_step if (epoch == start_epoch and start_step > 0) else 0
|
||||
batch_sampler = SkipBatchSampler(train_sampler or indices, args.batch_size, skip)
|
||||
loader = DataLoader(train_ds, batch_sampler=batch_sampler, num_workers=args.num_workers, pin_memory=True)
|
||||
loader = DataLoader(train_ds, batch_sampler=batch_sampler, num_workers=args.num_workers,
|
||||
pin_memory=should_pin_memory(args.device))
|
||||
if skip > 0:
|
||||
Logger(f'Epoch [{epoch + 1}/{args.epochs}]: 跳过前{start_step}个step，从step {start_step + 1}开始')
|
||||
train_epoch(epoch, loader, len(loader) + skip, start_step, wandb)
|
||||
|
|
@ -170,4 +172,4 @@ if __name__ == "__main__":
|
|||
# ========== 9. 清理分布进程 ==========
|
||||
if dist.is_initialized():
|
||||
dist.barrier()
|
||||
dist.destroy_process_group()
|
||||
dist.destroy_process_group()
|
||||
|
|
|
|||
|
|
@ -7,6 +7,7 @@ __package__ = "trainer"
|
|||
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
|
||||
import random
|
||||
import math
|
||||
from contextlib import nullcontext
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.distributed as dist
|
||||
|
|
@ -15,6 +16,91 @@ from torch.utils.data import Sampler
|
|||
from transformers import AutoTokenizer, AutoModel, AutoModelForSequenceClassification
|
||||
from model.model_minimind import MiniMindForCausalLM
|
||||
|
||||
|
||||
SUPPORTED_DTYPES = {
|
||||
'float32': torch.float32,
|
||||
'float16': torch.float16,
|
||||
'bfloat16': torch.bfloat16,
|
||||
}
|
||||
|
||||
|
||||
def get_device_type(device):
|
||||
"""Return the accelerator type for strings such as cuda:0, mps, and cpu."""
|
||||
return torch.device(device).type
|
||||
|
||||
|
||||
def get_default_device():
|
||||
"""Prefer CUDA, then Apple Silicon MPS, and finally CPU."""
|
||||
if torch.cuda.is_available():
|
||||
return 'cuda:0'
|
||||
if torch.backends.mps.is_available():
|
||||
return 'mps'
|
||||
return 'cpu'
|
||||
|
||||
|
||||
def validate_device(device):
|
||||
device_type = get_device_type(device)
|
||||
if device_type == 'cuda' and not torch.cuda.is_available():
|
||||
raise RuntimeError('CUDA was requested but is not available.')
|
||||
if device_type == 'mps' and not torch.backends.mps.is_available():
|
||||
reason = ('this PyTorch build has no MPS support'
|
||||
if not torch.backends.mps.is_built()
|
||||
else 'the current macOS device does not expose MPS')
|
||||
raise RuntimeError(f'MPS was requested but is not available: {reason}.')
|
||||
if device_type not in {'cuda', 'mps', 'cpu'}:
|
||||
raise ValueError(f'Unsupported training device: {device}')
|
||||
return device_type
|
||||
|
||||
|
||||
def get_torch_dtype(dtype):
|
||||
try:
|
||||
return SUPPORTED_DTYPES[dtype]
|
||||
except KeyError as exc:
|
||||
choices = ', '.join(SUPPORTED_DTYPES)
|
||||
raise ValueError(f'Unsupported dtype {dtype!r}; choose one of: {choices}') from exc
|
||||
|
||||
|
||||
def get_autocast_context(device, dtype):
|
||||
"""Create the appropriate AMP context for CUDA or Apple MPS."""
|
||||
device_type = validate_device(device)
|
||||
torch_dtype = get_torch_dtype(dtype)
|
||||
if device_type == 'cpu' or torch_dtype == torch.float32:
|
||||
return nullcontext()
|
||||
return torch.autocast(device_type=device_type, dtype=torch_dtype)
|
||||
|
||||
|
||||
def get_grad_scaler(device, dtype):
|
||||
"""Enable gradient scaling for fp16 CUDA/MPS and use a no-op otherwise."""
|
||||
device_type = validate_device(device)
|
||||
enabled = dtype == 'float16' and device_type in {'cuda', 'mps'}
|
||||
try:
|
||||
return torch.amp.GradScaler(device_type, enabled=enabled)
|
||||
except (TypeError, RuntimeError):
|
||||
# Older PyTorch releases only expose a CUDA GradScaler. Keep CUDA
|
||||
# compatibility and fail clearly for MPS instead of silently training
|
||||
# fp16 without scaling.
|
||||
if device_type == 'cuda':
|
||||
return torch.cuda.amp.GradScaler(enabled=enabled)
|
||||
if enabled:
|
||||
raise RuntimeError(
|
||||
'This PyTorch version cannot scale MPS float16 gradients. '
|
||||
'Upgrade PyTorch or use --dtype bfloat16/float32.'
|
||||
)
|
||||
return torch.cuda.amp.GradScaler(enabled=False)
|
||||
|
||||
|
||||
def should_pin_memory(device):
|
||||
"""Pinned host memory accelerates CUDA copies but does not benefit MPS."""
|
||||
return get_device_type(device) == 'cuda'
|
||||
|
||||
|
||||
def empty_device_cache(device):
|
||||
device_type = get_device_type(device)
|
||||
if device_type == 'cuda' and torch.cuda.is_available():
|
||||
torch.cuda.empty_cache()
|
||||
elif device_type == 'mps' and torch.backends.mps.is_available():
|
||||
torch.mps.empty_cache()
|
||||
|
||||
def get_model_params(model, config):
|
||||
total = sum(p.numel() for p in model.parameters()) / 1e6
|
||||
n_routed = getattr(config, 'n_routed_experts', getattr(config, 'num_experts', 0))
|
||||
|
|
@ -41,10 +127,17 @@ def get_lr(current_step, total_steps, lr):
|
|||
return lr*(0.1 + 0.45*(1 + math.cos(math.pi * current_step / total_steps)))
|
||||
|
||||
|
||||
def init_distributed_mode():
|
||||
def init_distributed_mode(device=None):
|
||||
if int(os.environ.get("RANK", -1)) == -1:
|
||||
return 0 # 非DDP模式
|
||||
|
||||
device = device or get_default_device()
|
||||
if get_device_type(device) != 'cuda':
|
||||
raise RuntimeError(
|
||||
'Distributed training in MiniMind currently requires CUDA/NCCL. '
|
||||
'Run MPS training as a single process with `python train_xxx.py --device mps`.'
|
||||
)
|
||||
|
||||
dist.init_process_group(backend="nccl")
|
||||
local_rank = int(os.environ["LOCAL_RANK"])
|
||||
torch.cuda.set_device(local_rank)
|
||||
|
|
@ -55,10 +148,13 @@ def setup_seed(seed: int):
|
|||
random.seed(seed)
|
||||
np.random.seed(seed)
|
||||
torch.manual_seed(seed)
|
||||
torch.cuda.manual_seed(seed)
|
||||
torch.cuda.manual_seed_all(seed)
|
||||
torch.backends.cudnn.deterministic = True
|
||||
torch.backends.cudnn.benchmark = False
|
||||
if torch.cuda.is_available():
|
||||
torch.cuda.manual_seed(seed)
|
||||
torch.cuda.manual_seed_all(seed)
|
||||
torch.backends.cudnn.deterministic = True
|
||||
torch.backends.cudnn.benchmark = False
|
||||
if torch.backends.mps.is_available():
|
||||
torch.mps.manual_seed(seed)
|
||||
|
||||
def lm_checkpoint(lm_config, weight='full_sft', model=None, optimizer=None, epoch=0, step=0, wandb=None, save_dir='../checkpoints', **kwargs):
|
||||
os.makedirs(save_dir, exist_ok=True)
|
||||
|
|
@ -69,6 +165,7 @@ def lm_checkpoint(lm_config, weight='full_sft', model=None, optimizer=None, epoc
|
|||
if model is not None:
|
||||
raw_model = model.module if isinstance(model, DistributedDataParallel) else model
|
||||
raw_model = getattr(raw_model, '_orig_mod', raw_model)
|
||||
model_device = next(raw_model.parameters()).device
|
||||
state_dict = raw_model.state_dict()
|
||||
state_dict = {k: v.half().cpu() for k, v in state_dict.items()}
|
||||
ckp_tmp = ckp_path + '.tmp'
|
||||
|
|
@ -103,7 +200,7 @@ def lm_checkpoint(lm_config, weight='full_sft', model=None, optimizer=None, epoc
|
|||
torch.save(resume_data, resume_tmp)
|
||||
os.replace(resume_tmp, resume_path)
|
||||
del state_dict, resume_data
|
||||
torch.cuda.empty_cache()
|
||||
empty_device_cache(model_device)
|
||||
else: # 加载模式
|
||||
if os.path.exists(resume_path):
|
||||
ckp_data = torch.load(resume_path, map_location='cpu')
|
||||
|
|
@ -116,7 +213,8 @@ def lm_checkpoint(lm_config, weight='full_sft', model=None, optimizer=None, epoc
|
|||
return None
|
||||
|
||||
|
||||
def init_model(lm_config, from_weight='pretrain', tokenizer_path='../model', save_dir='../out', device='cuda'):
|
||||
def init_model(lm_config, from_weight='pretrain', tokenizer_path='../model', save_dir='../out', device=None):
|
||||
device = device or get_default_device()
|
||||
tokenizer = AutoTokenizer.from_pretrained(tokenizer_path)
|
||||
model = MiniMindForCausalLM(lm_config)
|
||||
|
||||
|
|
@ -158,7 +256,8 @@ class SkipBatchSampler(Sampler):
|
|||
|
||||
|
||||
class LMForRewardModel:
|
||||
def __init__(self, model_path, device="cuda", dtype=torch.float16):
|
||||
def __init__(self, model_path, device=None, dtype=torch.float16):
|
||||
device = device or get_default_device()
|
||||
self.tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
|
||||
self.model = AutoModel.from_pretrained(model_path, torch_dtype=dtype, trust_remote_code=True)
|
||||
self.model = self.model.to(device).eval()
|
||||
|
|
@ -174,4 +273,4 @@ class LMForRewardModel:
|
|||
{"role": "assistant", "content": response}
|
||||
]
|
||||
score = self.model.get_score(self.tokenizer, eval_messages)
|
||||
return max(min(score, 3.0), -3.0)
|
||||
return max(min(score, 3.0), -3.0)
|
||||
|
|
|
|||
Loading…
Reference in New Issue