test: cover MPS device and training paths
mengxin
parent
fce15c0bf4
commit
abfd81222a
|
|
@ -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()
|
||||
Loading…
Reference in New Issue