用途:在本地 GPU 上运行 Unsloth 训练,并作为训练管理器负责环境校验、GPU 选择、进度监控与检查点处理。
主要能力:
- 前置检查:用 torch.cuda.is_available()、设备名与显存信息验证 CUDA;不可用时提示检查 NVIDIA 驱动(nvidia-smi)、CUDA(nvcc --version)并重装 PyTorch:pip install torch --index-url https://download.pytorch.org/whl/cu121。
- 依赖安装:pip install unsloth torch transformers trl peft datasets accelerate bitsandbytes。
- 显存参考:8GB 对应 7B/4bit/batch=1/LoRA r=8;12GB 对应 7B/4bit/batch=2/r=16;16GB 对应 7-13B/4bit/batch=2/r=16-32;24GB 对应 7-14B/4bit/batch=4/r=32。
- Docker 方案:docker run -d -e JUPYTER_PASSWORD="unsloth" -p 8888:8888 -v $(pwd)/work:/workspace/work --gpus all unsloth/unsloth,Jupyter 地址 http://localhost:8888;环境变量包括 JUPYTER_PASSWORD、JUPYTER_PORT、USER_PASSWORD。
- 运行训练:Notebook 方式 jupyter notebook notebooks/sft_template.ipynb;脚本方式先改配置再 python scripts/train_sft.py。
- 多 GPU 选择:设置 os.environ["CUDA_VISIBLE_DEVICES"]。
- 监控:watch -n 1 nvidia-smi 或 pip install nvitop && nvitop;可选设置 WANDB_API_KEY 并在 TrainingArguments 加 report_to="wandb"。
- 检查点:resume_from_checkpoint=True 自动找最新,或指定路径;保存 LoRA 权重与合并模型,并可用 FastLanguageModel 测试推理。
限制与排错:OOM 可依次降低 batch_size、增大 gradient_accumulation、减小 max_seq_length 与 LoRA rank,并调用 torch.cuda.empty_cache();损失不降时检查学习率、聊天模板与数据格式;过慢可启用 bf16 或 packing=True。训练前关闭其他 GPU 应用,温度保持 85C 以下,长时间运行建议 UPS。具体平台的收费与价格信息暂未提供。