1. Chuẩn bị Dataset và Dockerfile cho bài toán Training
Bước đầu tiên là chuẩn bị dữ liệu huấn luyện và đóng gói môi trường runtime vào Docker image để đảm bảo tính nhất quán trên mọi node.
Tạo thư mục làm việc và cấu trúc dữ liệu mẫu tại thư mục `training-data` trong container hoặc trên host.
mkdir -p ~/ai-training-project/data
cd ~/ai-training-project/data
# Tạo dataset giả lập dạng text để test Ray Data
python3 -c "import json; [json.dump({'text': f'Sample text data for training index {i}', 'label': i % 2}) for i in range(1000)]" > train_data.jsonl
Kết quả mong đợi: File `train_data.jsonl` được tạo với 1000 dòng dữ liệu JSON.
Viết file `Dockerfile` để đóng gói môi trường Python với Ray, DeepSpeed, PyTorch và các thư viện cần thiết.
Đường dẫn: `~/ai-training-project/Dockerfile`
FROM nvcr.io/nvidia/pytorch:23.10-py3
WORKDIR /app
# Cài đặt Ray và DeepSpeed
RUN pip install --upgrade pip && \
pip install ray[default] deepspeed transformers datasets torchmetrics && \
pip install accelerate
# Copy code và data vào container
COPY training_script.py .
COPY data/train_data.jsonl ./data/
# Set môi trường
ENV RAY_ADDRESS="http://ray-head:8265"
CMD ["python", "training_script.py"]
Kết quả mong đợi: File Dockerfile sẵn sàng để build image. Đảm bảo bạn đã build và push image lên registry nội bộ (ví dụ: `local-registry:5000/ai-training:v1`) ở bước trước.
Build và push Docker image lên registry của cluster Kubernetes.
cd ~/ai-training-project
docker build -t local-registry:5000/ai-training:v1 .
docker push local-registry:5000/ai-training:v1
Kết quả mong đợi: Image được đẩy thành công lên registry. Bạn có thể verify bằng lệnh `docker images` trên worker node.
2. Viết Script Training tích hợp DeepSpeed ZeRO và Ray Data
Viết script Python chính để xử lý dữ liệu bằng Ray Data, phân phối training bằng Ray Train và tối ưu hóa bộ nhớ bằng DeepSpeed ZeRO Stage 2.
File script này sẽ thực hiện: (1) Khởi tạo Ray Client, (2) Chuyển đổi dữ liệu sang Ray Dataset, (3) Cấu hình DeepSpeedConfig, (4) Chạy training loop.
Đường dẫn: `~/ai-training-project/training_script.py`
import os
import json
import ray
from ray import train
from ray.train import ScalingConfig, RunConfig
from ray.train.torch import TorchTrainer
from ray.data import Dataset
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch
from torch.utils.data import Dataset as TorchDataset, DataLoader
from deepspeed import DeepSpeedConfig
# Cấu hình môi trường
MODEL_NAME = "facebook/opt-350m"
BATCH_SIZE = 4
EPOCHS = 2
class TextDataset(TorchDataset):
def __init__(self, data):
self.data = data
self.tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
def __len__(self):
return len(self.data)
def __getitem__(self, idx):
item = self.data[idx]
inputs = self.tokenizer(item['text'], return_tensors="pt", padding=True, truncation=True, max_length=128)
return {
'input_ids': inputs['input_ids'].squeeze(0),
'attention_mask': inputs['attention_mask'].squeeze(0),
'labels': torch.tensor(item['label'])
}
def load_data():
# Sử dụng Ray Data để đọc file JSONL
ds = ray.data.read_json("data/train_data.jsonl")
return ds
def train_func():
# Load model và tokenizer
model = AutoModelForCausalLM.from_pretrained(MODEL_NAME)
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
# Load data từ Ray Data
dataset = load_data()
# Chuyển đổi sang torch dataset cho training
torch_ds = TextDataset(dataset.to_torch())
dataloader = DataLoader(torch_ds, batch_size=BATCH_SIZE, shuffle=True)
# Cấu hình DeepSpeed ZeRO Stage 2
ds_config = {
"train_batch_size": BATCH_SIZE,
"train_micro_batch_size_per_gpu": BATCH_SIZE,
"gradient_accumulation_steps": 1,
"zero_optimization": {
"stage": 2,
"offload_optimizer": {"device": "cpu"},
"offload_param": {"device": "cpu"}
},
"fp16": {"enabled": True}
}
# Khởi tạo Trainer với DeepSpeedConfig
trainer = TorchTrainer(
train_func_inner=model,
train_loop_dir=lambda model, dataloader: train_loop(model, dataloader, tokenizer),
scaling_config=ScalingConfig(num_workers=4, use_gpu=True),
run_config=RunConfig(name="distributed-opt-training"),
train_loop_config=ds_config
)
# Lưu checkpoint
trainer.fit()
def train_loop(model, dataloader, tokenizer):
model.train()
optimizer = torch.optim.AdamW(model.parameters(), lr=5e-5)
for epoch in range(EPOCHS):
total_loss = 0
for batch in dataloader:
inputs = batch['input_ids']
attention_mask = batch['attention_mask']
labels = batch['labels']
outputs = model(
input_ids=inputs,
attention_mask=attention_mask,
labels=labels
)
loss = outputs.loss
loss.backward()
optimizer.step()
optimizer.zero_grad()
total_loss += loss.item()
# Report metrics to Ray Dashboard
train.report({"loss": total_loss / len(dataloader), "epoch": epoch})
def train_func_inner(model):
# Hàm wrapper để Ray gọi
pass
if __name__ == "__main__":
# Khởi tạo Ray Client (kết nối vào Ray Cluster đã deploy)
ray.init(address=os.environ.get("RAY_ADDRESS", "ray://ray-head:8265"))
# Chạy training
train_func()
print("Training completed successfully!")
ray.shutdown()
Kết quả mong đợi: Script được lưu thành file `.py`. Logic đã tích hợp sẵn Ray Data để phân phối dữ liệu và DeepSpeed ZeRO Stage 2 để giảm bộ nhớ GPU.
Verify script bằng cách chạy kiểm tra cú pháp Python.
python3 -m py_compile ~/ai-training-project/training_script.py
Kết quả mong đợi: Không có lỗi (no error message). Nếu có lỗi, hãy kiểm tra lại import và logic.
3. Chạy thử nghiệm trên Cluster và Theo dõi qua Ray Dashboard
Tạo Kubernetes Job để triển khai Ray Cluster và chạy task training song song trên các worker GPU.
Tạo file manifest Kubernetes `ray-job.yaml` để định nghĩa Ray Cluster (Head + Workers) và Ray Job.
Đường dẫn: `~/ai-training-project/ray-job.yaml`
apiVersion: v1
kind: Namespace
metadata:
name: ai-training
---
apiVersion: ray.io/v1
kind: RayCluster
metadata:
name: opt-training-cluster
namespace: ai-training
spec:
headGroupSpec:
rayStartParams:
dashboard-host: 0.0.0.0
num-cpus: "2"
template:
spec:
containers:
- name: ray-head
image: local-registry:5000/ai-training:v1
ports:
- containerPort: 6379
name: gcs-server
- containerPort: 8265
name: dashboard
- containerPort: 10001
name: client
resources:
limits:
cpu: "2"
memory: 4Gi
env:
- name: RAY_ADDRESS
value: "http://opt-training-cluster-head-svc:8265"
volumeMounts:
- name: data
mountPath: /app/data
volumes:
- name: data
hostPath:
path: /mnt/data/train_data
type: Directory
workerGroupSpecs:
- groupName: gpu-workers
replicas: 4
rayStartParams:
num-cpus: "4"
num-gpus: "1"
template:
spec:
containers:
- name: ray-worker
image: local-registry:5000/ai-training:v1
resources:
limits:
cpu: "4"
memory: 16Gi
nvidia.com/gpu: "1"
env:
- name: RAY_ADDRESS
value: "http://opt-training-cluster-head-svc:8265"
volumeMounts:
- name: data
mountPath: /app/data
volumes:
- name: data
hostPath:
path: /mnt/data/train_data
type: Directory
---
apiVersion: ray.io/v1
kind: RayJob
metadata:
name: opt-training-job
namespace: ai-training
spec:
clusterSelector:
name: opt-training-cluster
entrypoint: "python training_script.py"
runtimeEnv:
workingDir: "/app"
envVars:
RAY_ADDRESS: "http://opt-training-cluster-head-svc:8265"
Kết quả mong đợi: File YAML được tạo với cấu hình Ray Cluster có 1 Head Node và 4 Worker Nodes (mỗi node 1 GPU), cùng với Ray Job để chạy script.
Triển khai cluster và job lên Kubernetes.
cd ~/ai-training-project
kubectl apply -f ray-job.yaml
Kết quả mong đợi: Kubernetes tạo ra các Pod `opt-training-cluster-head-...` và `opt-training-cluster-gpu-workers-...`. Trạng thái chuyển dần từ `Pending` sang `Running`.
Đợi Pod khởi động xong và expose Ray Dashboard để theo dõi tiến độ.
kubectl port-forward svc/opt-training-cluster-head-svc -n ai-training 8265:8265
Kết quả mong đợi: Terminal hiển thị dòng "Forwarding from 127.0.0.1:8265 -> 8265". Mở trình duyệt truy cập `http://localhost:8265`.
Kiểm tra trạng thái của Ray Job.
kubectl get rayjob opt-training-job -n ai-training
Kết quả mong đợi: Trạng thái `status.jobStatus` chuyển từ `PENDING` -> `SUBMITTED` -> `RUNNING` -> `SUCCEEDED`.
Verify tiến độ training qua Ray Dashboard.
Mở tab "Dashboard" -> chọn job "opt-training-job". Bạn sẽ thấy:
- Tab "Cluster": Hiển thị 5 node (1 head + 4 workers) đang hoạt động, GPU utilization tăng cao.
- Tab "Metrics": Biểu đồ `loss` giảm dần theo epoch.
- Tab "Logs": Output console của script training.
Kết quả mong đợi: Biểu đồ loss hội tụ và job chuyển trạng thái thành `SUCCEEDED` sau khi hoàn thành 2 epoch.
Verify log chi tiết từ container để đảm bảo không có lỗi DeepSpeed hoặc Ray.
kubectl logs -f rayjob/opt-training-job -n ai-training --selector ray.io/cluster=opt-training-cluster -c ray-worker | grep "Training completed"
Kết quả mong đợi: Xuất hiện dòng "Training completed successfully!" trong log.
4. Xử lý sự cố và Kiểm tra tài nguyên
Trong quá trình chạy, nếu gặp lỗi OOM (Out of Memory) hoặc lỗi GPU, hãy kiểm tra cấu hình DeepSpeed.
Chạy lệnh để xem GPU utilization thực tế trên các node worker.
kubectl get pods -n ai-training | grep worker
# Sau đó ssh vào 1 node worker để chạy nvidia-smi
ssh node-worker-1 "nvidia-smi"
Kết quả mong đợi: `nvidia-smi` hiển thị GPU đang được sử dụng bởi quá trình `python` với VRAM usage cao nhưng không bị lỗi.
Nếu job bị lỗi `OOMKilled`, giảm `train_micro_batch_size_per_gpu` trong `ds_config` trong file `training_script.py` xuống 2 hoặc 1, sau đó rebuild image và chạy lại.
Xóa tài nguyên sau khi test xong để giải phóng GPU.
kubectl delete -f ~/ai-training-project/ray-job.yaml
Kết quả mong đợi: Cluster và Job bị xóa, Pod quay về trạng thái không còn, GPU được giải phóng cho các task khác.
Điều hướng series:
Mục lục: Series: Series: Xây dựng nền tảng AI Training phân tán với DeepSpeed, Ray và Kubernetes trên hạ tầng Proxmox
« Phần 4: Triển khai Ray và DeepSpeed trên Kubernetes
Phần 6: Tối ưu hóa hiệu năng và cân bằng tải cho training »