mirror of
https://github.com/hanzo-ml/trainer.git
synced 2026-09-24 00:31:34 +00:00
* fix(examples): Verify TrainJob Completion Signed-off-by: Andrey Velichkevich <[email protected]> * Rename client to TrainerClient() Signed-off-by: Andrey Velichkevich <[email protected]> --------- Signed-off-by: Andrey Velichkevich <[email protected]>
29 KiB
29 KiB
In [ ]:
# !pip install -U kubeflowIn [ ]:
!pip install torch==2.10.0
!pip install torchvision==0.22.1In [1]:
def train_fashion_mnist(num_epochs=2, batch_size=100, lr=0.1, momentum=0.9):
import os
import torch
import torch.distributed as dist
import torch.nn.functional as F
from torch import nn
from torch.utils.data import DataLoader, DistributedSampler
from torchvision import datasets, transforms
# Define the PyTorch CNN model to be trained
class Net(nn.Module):
def __init__(self):
super(Net, self).__init__()
self.conv1 = nn.Conv2d(1, 20, 5, 1)
self.conv2 = nn.Conv2d(20, 50, 5, 1)
self.fc1 = nn.Linear(4 * 4 * 50, 500)
self.fc2 = nn.Linear(500, 10)
def forward(self, x):
x = F.relu(self.conv1(x))
x = F.max_pool2d(x, 2, 2)
x = F.relu(self.conv2(x))
x = F.max_pool2d(x, 2, 2)
x = x.view(-1, 4 * 4 * 50)
x = F.relu(self.fc1(x))
x = self.fc2(x)
return F.log_softmax(x, dim=1)
# Use NCCL if a GPU is available, otherwise use Gloo as communication backend.
device, backend = ("cuda", "nccl") if torch.cuda.is_available() else ("cpu", "gloo")
print(f"Using Device: {device}, Backend: {backend}")
# Setup PyTorch distributed.
local_rank = int(os.getenv("LOCAL_RANK", 0))
dist.init_process_group(backend=backend)
print(
"Distributed Training for WORLD_SIZE: {}, RANK: {}, LOCAL_RANK: {}".format(
dist.get_world_size(),
dist.get_rank(),
local_rank,
)
)
# Create the model and load it into the device.
device = torch.device(f"{device}:{local_rank}")
model =Net().to(device)
if torch.cuda.is_available():
model=torch.compile(model)
model = nn.parallel.DistributedDataParallel(model)
optimizer = torch.optim.SGD(model.parameters(), lr=lr, momentum=momentum)
# Download FashionMNIST dataset only on local_rank=0 process.
if local_rank == 0:
dataset = datasets.FashionMNIST(
"./data",
train=True,
download=True,
transform=transforms.Compose([transforms.ToTensor()]),
)
dist.barrier()
dataset = datasets.FashionMNIST(
"./data",
train=True,
download=False,
transform=transforms.Compose([transforms.ToTensor()]),
)
# Shard the dataset across workers.
train_loader = DataLoader(
dataset,
batch_size=batch_size,
sampler=DistributedSampler(dataset)
)
dist.barrier()
for epoch in range(1, num_epochs + 1):
model.train()
# Iterate over mini-batches from the training set
for batch_idx, (inputs, labels) in enumerate(train_loader):
# Copy the data to the GPU device if available
inputs, labels = inputs.to(device), labels.to(device)
# Forward pass
outputs = model(inputs)
loss = F.nll_loss(outputs, labels)
# Backward pass
optimizer.zero_grad()
loss.backward()
optimizer.step()
if batch_idx % 10 == 0 and dist.get_rank() == 0:
print(
"Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}".format(
epoch,
batch_idx * len(inputs),
len(train_loader.dataset),
100.0 * batch_idx / len(train_loader),
loss.item(),
)
)
# Wait for the distributed training to complete
dist.barrier()
if dist.get_rank() == 0:
print("Training is finished")
# Finally clean up PyTorch distributed
dist.destroy_process_group()In [3]:
from kubeflow.trainer import CustomTrainer, TrainerClient, LocalProcessBackendConfig
# Initialize local backend
backend_config = LocalProcessBackendConfig(cleanup_venv=True)
client = TrainerClient(backend_config=backend_config)
# List available runtimes
for runtime in client.list_runtimes():
if runtime.name == "torch-distributed":
torch_runtime = runtime
break
# Submit training job
job_name = client.train(
trainer=CustomTrainer(
func=train_fashion_mnist,
packages_to_install=["torch", "torchvision"],
),
runtime=torch_runtime,
)
# Stream logs
for logline in client.get_job_logs(job_name, follow=True):
print(logline, end='')Using Device: cpu, Backend: gloo Distributed Training for WORLD_SIZE: 1, RANK: 0, LOCAL_RANK: 0 Train Epoch: 1 [0/60000 (0%)] Loss: 2.300957 Train Epoch: 1 [1000/60000 (2%)] Loss: 2.068733 Train Epoch: 1 [2000/60000 (3%)] Loss: 1.136209 Train Epoch: 1 [3000/60000 (5%)] Loss: 0.897645 Train Epoch: 1 [4000/60000 (7%)] Loss: 0.868213 Train Epoch: 1 [5000/60000 (8%)] Loss: 0.943094 Train Epoch: 1 [6000/60000 (10%)] Loss: 0.785388 Train Epoch: 1 [7000/60000 (12%)] Loss: 0.772929 Train Epoch: 1 [8000/60000 (13%)] Loss: 0.612156 Train Epoch: 1 [9000/60000 (15%)] Loss: 0.783617 Train Epoch: 1 [10000/60000 (17%)] Loss: 0.608218 Train Epoch: 1 [11000/60000 (18%)] Loss: 0.595062 Train Epoch: 1 [12000/60000 (20%)] Loss: 0.612815 Train Epoch: 1 [13000/60000 (22%)] Loss: 0.523903 Train Epoch: 1 [14000/60000 (23%)] Loss: 0.673269 Train Epoch: 1 [15000/60000 (25%)] Loss: 0.675361 Train Epoch: 1 [16000/60000 (27%)] Loss: 0.591610 Train Epoch: 1 [17000/60000 (28%)] Loss: 0.515040 Train Epoch: 1 [18000/60000 (30%)] Loss: 0.535079 Train Epoch: 1 [19000/60000 (32%)] Loss: 0.721351 Train Epoch: 1 [20000/60000 (33%)] Loss: 0.413175 Train Epoch: 1 [21000/60000 (35%)] Loss: 0.407277 Train Epoch: 1 [22000/60000 (37%)] Loss: 0.356355 Train Epoch: 1 [23000/60000 (38%)] Loss: 0.685514 Train Epoch: 1 [24000/60000 (40%)] Loss: 0.439642 Train Epoch: 1 [25000/60000 (42%)] Loss: 0.399821 Train Epoch: 1 [26000/60000 (43%)] Loss: 0.556694 Train Epoch: 1 [27000/60000 (45%)] Loss: 0.326490 Train Epoch: 1 [28000/60000 (47%)] Loss: 0.401890 Train Epoch: 1 [29000/60000 (48%)] Loss: 0.449772 Train Epoch: 1 [30000/60000 (50%)] Loss: 0.427323 Train Epoch: 1 [31000/60000 (52%)] Loss: 0.400165 Train Epoch: 1 [32000/60000 (53%)] Loss: 0.334357 Train Epoch: 1 [33000/60000 (55%)] Loss: 0.298294 Train Epoch: 1 [34000/60000 (57%)] Loss: 0.551655 Train Epoch: 1 [35000/60000 (58%)] Loss: 0.456408 Train Epoch: 1 [36000/60000 (60%)] Loss: 0.409366 Train Epoch: 1 [37000/60000 (62%)] Loss: 0.449682 Train Epoch: 1 [38000/60000 (63%)] Loss: 0.433943 Train Epoch: 1 [39000/60000 (65%)] Loss: 0.403032 Train Epoch: 1 [40000/60000 (67%)] Loss: 0.427653 Train Epoch: 1 [41000/60000 (68%)] Loss: 0.596136 Train Epoch: 1 [42000/60000 (70%)] Loss: 0.451446 Train Epoch: 1 [43000/60000 (72%)] Loss: 0.415507 Train Epoch: 1 [44000/60000 (73%)] Loss: 0.431509 Train Epoch: 1 [45000/60000 (75%)] Loss: 0.269870 Train Epoch: 1 [46000/60000 (77%)] Loss: 0.603037 Train Epoch: 1 [47000/60000 (78%)] Loss: 0.361142 Train Epoch: 1 [48000/60000 (80%)] Loss: 0.533713 Train Epoch: 1 [49000/60000 (82%)] Loss: 0.334492 Train Epoch: 1 [50000/60000 (83%)] Loss: 0.360594 Train Epoch: 1 [51000/60000 (85%)] Loss: 0.479819 Train Epoch: 1 [52000/60000 (87%)] Loss: 0.408797 Train Epoch: 1 [53000/60000 (88%)] Loss: 0.402639 Train Epoch: 1 [54000/60000 (90%)] Loss: 0.291429 Train Epoch: 1 [55000/60000 (92%)] Loss: 0.308326 Train Epoch: 1 [56000/60000 (93%)] Loss: 0.503240 Train Epoch: 1 [57000/60000 (95%)] Loss: 0.449824 Train Epoch: 1 [58000/60000 (97%)] Loss: 0.339513 Train Epoch: 1 [59000/60000 (98%)] Loss: 0.260989 Train Epoch: 2 [0/60000 (0%)] Loss: 0.379833 Train Epoch: 2 [1000/60000 (2%)] Loss: 0.421474 Train Epoch: 2 [2000/60000 (3%)] Loss: 0.298147 Train Epoch: 2 [3000/60000 (5%)] Loss: 0.302185 Train Epoch: 2 [4000/60000 (7%)] Loss: 0.407515 Train Epoch: 2 [5000/60000 (8%)] Loss: 0.372201 Train Epoch: 2 [6000/60000 (10%)] Loss: 0.403835 Train Epoch: 2 [7000/60000 (12%)] Loss: 0.403082 Train Epoch: 2 [8000/60000 (13%)] Loss: 0.429058 Train Epoch: 2 [9000/60000 (15%)] Loss: 0.427616 Train Epoch: 2 [10000/60000 (17%)] Loss: 0.353761 Train Epoch: 2 [11000/60000 (18%)] Loss: 0.392035 Train Epoch: 2 [12000/60000 (20%)] Loss: 0.315083 Train Epoch: 2 [13000/60000 (22%)] Loss: 0.303653 Train Epoch: 2 [14000/60000 (23%)] Loss: 0.477725 Train Epoch: 2 [15000/60000 (25%)] Loss: 0.311283 Train Epoch: 2 [16000/60000 (27%)] Loss: 0.320934 Train Epoch: 2 [17000/60000 (28%)] Loss: 0.344658 Train Epoch: 2 [18000/60000 (30%)] Loss: 0.252863 Train Epoch: 2 [19000/60000 (32%)] Loss: 0.433338 Train Epoch: 2 [20000/60000 (33%)] Loss: 0.401821 Train Epoch: 2 [21000/60000 (35%)] Loss: 0.277403 Train Epoch: 2 [22000/60000 (37%)] Loss: 0.324770 Train Epoch: 2 [23000/60000 (38%)] Loss: 0.450259 Train Epoch: 2 [24000/60000 (40%)] Loss: 0.280156 Train Epoch: 2 [25000/60000 (42%)] Loss: 0.273340 Train Epoch: 2 [26000/60000 (43%)] Loss: 0.405919 Train Epoch: 2 [27000/60000 (45%)] Loss: 0.270347 Train Epoch: 2 [28000/60000 (47%)] Loss: 0.389775 Train Epoch: 2 [29000/60000 (48%)] Loss: 0.266520 Train Epoch: 2 [30000/60000 (50%)] Loss: 0.369550 Train Epoch: 2 [31000/60000 (52%)] Loss: 0.307309 Train Epoch: 2 [32000/60000 (53%)] Loss: 0.282144 Train Epoch: 2 [33000/60000 (55%)] Loss: 0.259714 Train Epoch: 2 [34000/60000 (57%)] Loss: 0.455187 Train Epoch: 2 [35000/60000 (58%)] Loss: 0.418436 Train Epoch: 2 [36000/60000 (60%)] Loss: 0.347754 Train Epoch: 2 [37000/60000 (62%)] Loss: 0.285930 Train Epoch: 2 [38000/60000 (63%)] Loss: 0.298169 Train Epoch: 2 [39000/60000 (65%)] Loss: 0.315216 Train Epoch: 2 [40000/60000 (67%)] Loss: 0.207120 Train Epoch: 2 [41000/60000 (68%)] Loss: 0.447531 Train Epoch: 2 [42000/60000 (70%)] Loss: 0.356669 Train Epoch: 2 [43000/60000 (72%)] Loss: 0.295995 Train Epoch: 2 [44000/60000 (73%)] Loss: 0.306820 Train Epoch: 2 [45000/60000 (75%)] Loss: 0.209821 Train Epoch: 2 [46000/60000 (77%)] Loss: 0.474326 Train Epoch: 2 [47000/60000 (78%)] Loss: 0.292024 Train Epoch: 2 [48000/60000 (80%)] Loss: 0.554571 Train Epoch: 2 [49000/60000 (82%)] Loss: 0.313117 Train Epoch: 2 [50000/60000 (83%)] Loss: 0.247606 Train Epoch: 2 [51000/60000 (85%)] Loss: 0.385026 Train Epoch: 2 [52000/60000 (87%)] Loss: 0.369550 Train Epoch: 2 [53000/60000 (88%)] Loss: 0.326454 Train Epoch: 2 [54000/60000 (90%)] Loss: 0.317689 Train Epoch: 2 [55000/60000 (92%)] Loss: 0.296784 Train Epoch: 2 [56000/60000 (93%)] Loss: 0.399627 Train Epoch: 2 [57000/60000 (95%)] Loss: 0.408119 Train Epoch: 2 [58000/60000 (97%)] Loss: 0.341115 Train Epoch: 2 [59000/60000 (98%)] Loss: 0.191851 Training is finished
In [2]:
from kubeflow.trainer import CustomTrainer, TrainerClient
client = TrainerClient()In [3]:
for runtime in client.list_runtimes():
print(runtime)
if runtime.name == "torch-distributed":
torch_runtime = runtimeRuntime(name='deepspeed-distributed', trainer=RuntimeTrainer(trainer_type=<TrainerType.CUSTOM_TRAINER: 'CustomTrainer'>, framework='deepspeed', num_nodes=1, device='Unknown', device_count='1'), pretrained_model=None) Runtime(name='mlx-distributed', trainer=RuntimeTrainer(trainer_type=<TrainerType.CUSTOM_TRAINER: 'CustomTrainer'>, framework='mlx', num_nodes=1, device='gpu', device_count='1'), pretrained_model=None) Runtime(name='torch-distributed', trainer=RuntimeTrainer(trainer_type=<TrainerType.CUSTOM_TRAINER: 'CustomTrainer'>, framework='torch', num_nodes=1, device='Unknown', device_count='Unknown'), pretrained_model=None) Runtime(name='torchtune-llama3.2-1b', trainer=RuntimeTrainer(trainer_type=<TrainerType.BUILTIN_TRAINER: 'BuiltinTrainer'>, framework='torchtune', num_nodes=1, device='gpu', device_count='2'), pretrained_model=None) Runtime(name='torchtune-llama3.2-3b', trainer=RuntimeTrainer(trainer_type=<TrainerType.BUILTIN_TRAINER: 'BuiltinTrainer'>, framework='torchtune', num_nodes=1, device='gpu', device_count='2.0'), pretrained_model=None)
In [4]:
job_name = client.train(
trainer=CustomTrainer(
func=train_fashion_mnist,
func_args={"num_epochs": 3, "batch_size": 64, "lr": 0.01, "momentum": 0.9},
# Set how many PyTorch nodes you want to use for distributed training.
num_nodes=3,
# Set the resources for each PyTorch node.
resources_per_node={
"cpu": 3,
"memory": "16Gi",
# Uncomment this to distribute the TrainJob using GPU nodes.
# "nvidia.com/gpu": 1,
},
),
runtime=torch_runtime,
)In [5]:
# Wait for the running status.
client.wait_for_job_status(name=job_name, status={"Running"})Out [5]:
TrainJob(name='j819fe212f8e', creation_timestamp=datetime.datetime(2025, 9, 3, 13, 19, 56, tzinfo=TzInfo(UTC)), runtime=Runtime(name='torch-distributed', trainer=RuntimeTrainer(trainer_type=<TrainerType.CUSTOM_TRAINER: 'CustomTrainer'>, framework='torch', num_nodes=1, device='Unknown', device_count='Unknown'), pretrained_model=None), steps=[Step(name='node-0', status='Running', pod_name='j819fe212f8e-node-0-0-nlcqt', device='gpu', device_count='1'), Step(name='node-1', status='Running', pod_name='j819fe212f8e-node-0-1-rjztc', device='gpu', device_count='1'), Step(name='node-2', status='Running', pod_name='j819fe212f8e-node-0-2-l575c', device='gpu', device_count='1')], num_nodes=3, status='Running')
In [6]:
for c in client.get_job(name=job_name).steps:
print(f"Step: {c.name}, Status: {c.status}, Devices: {c.device} x {c.device_count}\n")Step: node-0, Status: Running, Devices: gpu x 1 Step: node-1, Status: Running, Devices: gpu x 1 Step: node-2, Status: Running, Devices: gpu x 1
In [7]:
for logline in client.get_job_logs(job_name, follow=True):
print(logline)[W903 13:20:06.250970596 socket.cpp:755] [c10d] The IPv6 network addresses of (j819fe212f8e-node-0-0.j819fe212f8e, 29500) cannot be retrieved (gai error: -2 - Name or service not known). [W903 13:20:07.936850072 socket.cpp:755] [c10d] The IPv6 network addresses of (j819fe212f8e-node-0-0.j819fe212f8e, 29500) cannot be retrieved (gai error: -2 - Name or service not known). [W903 13:20:08.873279869 socket.cpp:755] [c10d] The IPv6 network addresses of (j819fe212f8e-node-0-0.j819fe212f8e, 29500) cannot be retrieved (gai error: -2 - Name or service not known). [W903 13:20:08.673190939 socket.cpp:755] [c10d] The IPv6 network addresses of (j819fe212f8e-node-0-0.j819fe212f8e, 29500) cannot be retrieved (gai error: -2 - Name or service not known). [W903 13:20:09.683786251 socket.cpp:755] [c10d] The IPv6 network addresses of (j819fe212f8e-node-0-0.j819fe212f8e, 29500) cannot be retrieved (gai error: -2 - Name or service not known). Using Device: cuda, Backend: nccl Distributed Training for WORLD_SIZE: 3, RANK: 0, LOCAL_RANK: 0 100%|██████████| 26.4M/26.4M [00:02<00:00, 12.2MB/s] 100%|██████████| 29.5k/29.5k [00:00<00:00, 213kB/s] 100%|██████████| 4.42M/4.42M [00:01<00:00, 3.85MB/s] 100%|██████████| 5.15k/5.15k [00:00<00:00, 43.7MB/s] /opt/conda/lib/python3.11/site-packages/torch/distributed/distributed_c10d.py:4631: UserWarning: No device id is provided via `init_process_group` or `barrier `. Using the current device set by the user. warnings.warn( # warn only once Train Epoch: 1 [0/60000 (0%)] Loss: 2.297219 Train Epoch: 1 [1000/60000 (5%)] Loss: 1.803301 Train Epoch: 1 [2000/60000 (10%)] Loss: 1.663171 Train Epoch: 1 [3000/60000 (15%)] Loss: 1.462810 Train Epoch: 1 [4000/60000 (20%)] Loss: 1.270390 Train Epoch: 1 [5000/60000 (25%)] Loss: 0.977784 Train Epoch: 1 [6000/60000 (30%)] Loss: 0.791962 Train Epoch: 1 [7000/60000 (35%)] Loss: 0.626636 Train Epoch: 1 [8000/60000 (40%)] Loss: 0.613735 Train Epoch: 1 [9000/60000 (45%)] Loss: 0.590524 Train Epoch: 1 [10000/60000 (50%)] Loss: 0.529664 Train Epoch: 1 [11000/60000 (55%)] Loss: 0.448692 Train Epoch: 1 [12000/60000 (60%)] Loss: 0.508065 Train Epoch: 1 [13000/60000 (65%)] Loss: 0.519867 Train Epoch: 1 [14000/60000 (70%)] Loss: 0.485550 Train Epoch: 1 [15000/60000 (75%)] Loss: 0.377857 Train Epoch: 1 [16000/60000 (80%)] Loss: 0.456072 Train Epoch: 1 [17000/60000 (85%)] Loss: 0.563855 Train Epoch: 1 [18000/60000 (90%)] Loss: 0.426923 Train Epoch: 1 [19000/60000 (95%)] Loss: 0.379981 Train Epoch: 2 [0/60000 (0%)] Loss: 0.416116 Train Epoch: 2 [1000/60000 (5%)] Loss: 0.506344 Train Epoch: 2 [2000/60000 (10%)] Loss: 0.510181 Train Epoch: 2 [3000/60000 (15%)] Loss: 0.388432 Train Epoch: 2 [4000/60000 (20%)] Loss: 0.361425 Train Epoch: 2 [5000/60000 (25%)] Loss: 0.428971 Train Epoch: 2 [6000/60000 (30%)] Loss: 0.409087 Train Epoch: 2 [7000/60000 (35%)] Loss: 0.423781 Train Epoch: 2 [8000/60000 (40%)] Loss: 0.530143 Train Epoch: 2 [9000/60000 (45%)] Loss: 0.390188 Train Epoch: 2 [10000/60000 (50%)] Loss: 0.359034 Train Epoch: 2 [11000/60000 (55%)] Loss: 0.353203 Train Epoch: 2 [12000/60000 (60%)] Loss: 0.394956 Train Epoch: 2 [13000/60000 (65%)] Loss: 0.407927 Train Epoch: 2 [14000/60000 (70%)] Loss: 0.296730 Train Epoch: 2 [15000/60000 (75%)] Loss: 0.318915 Train Epoch: 2 [16000/60000 (80%)] Loss: 0.317562 Train Epoch: 2 [17000/60000 (85%)] Loss: 0.396108 Train Epoch: 2 [18000/60000 (90%)] Loss: 0.302800 Train Epoch: 2 [19000/60000 (95%)] Loss: 0.364170 Training is finished
In [ ]:
client.wait_for_job_status(name=job_name, timeout=20)In [ ]:
# client.delete_job(job_name)In [ ]: