diff --git a/tests/benchmark_mps.py b/tests/benchmark_mps.py new file mode 100644 index 0000000..c9037af --- /dev/null +++ b/tests/benchmark_mps.py @@ -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() diff --git a/tests/fixtures/pretrain_smoke.jsonl b/tests/fixtures/pretrain_smoke.jsonl new file mode 100644 index 0000000..743f55b --- /dev/null +++ b/tests/fixtures/pretrain_smoke.jsonl @@ -0,0 +1,8 @@ +{"text":"苹果芯片可以通过 MPS 加速 PyTorch 模型训练。"} +{"text":"MiniMind 是一个用于学习语言模型训练流程的小型项目。"} +{"text":"梯度下降通过反向传播更新神经网络参数。"} +{"text":"混合精度训练可以降低内存占用并提高吞吐量。"} +{"text":"统一内存允许 CPU 和 GPU 共享同一块物理内存。"} +{"text":"检查点可用于恢复模型参数和优化器状态。"} +{"text":"注意力机制根据上下文计算不同词元之间的关系。"} +{"text":"小规模冒烟测试可以快速发现训练环境问题。"} diff --git a/tests/fixtures/sft_smoke.jsonl b/tests/fixtures/sft_smoke.jsonl new file mode 100644 index 0000000..f39c67d --- /dev/null +++ b/tests/fixtures/sft_smoke.jsonl @@ -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":""}]} diff --git a/tests/test_device_utils.py b/tests/test_device_utils.py new file mode 100644 index 0000000..78f8184 --- /dev/null +++ b/tests/test_device_utils.py @@ -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()