# Volcano <-> Kubeflow Trainer Integration Configuration Sample # Reference: https://www.kubeflow.org/docs/components/trainer/gang-scheduling/volcano/ --- # 0. ConfigMap for test training script (VGG11 + CIFAR10 + NCCL) apiVersion: v1 kind: ConfigMap metadata: name: training-scripts data: train_nccl.py: | import os import time import torch import argparse import torch.distributed as dist import torch.multiprocessing as mp from torch.optim.lr_scheduler import StepLR # for dataset from torchvision.datasets.cifar import CIFAR10 import torchvision.transforms as tfs from torch.utils.data import DataLoader from torch.utils.data.distributed import DistributedSampler # for model from torchvision.models import vgg11 from torch.nn.parallel import DistributedDataParallel as DDP import numpy as np import random import datetime def set_random_seeds(random_seed=0): torch.manual_seed(random_seed) torch.backends.cudnn.deterministic = True torch.backends.cudnn.benchmark = False np.random.seed(random_seed) random.seed(random_seed) def get_args_parser(): parser = argparse.ArgumentParser(add_help=False) parser.add_argument('--lr', type=float, default=0.01) parser.add_argument('--epoch', type=int, default=90) parser.add_argument('--batch_size', type=int, default=1200) parser.add_argument('--global_rank', type=int, default=0) parser.add_argument('--vis_step', type=int, default=10) parser.add_argument('--num_workers', type=int, default=24) parser.add_argument("--local_rank", type=int, help="Local rank. Necessary for using the torch.distributed.launch utility.") parser.add_argument('--world_size', type=int, default=0) parser.add_argument('--port', type=int, default=2022) parser.add_argument('--root', type=str, default='data') parser.add_argument('--start_epoch', type=int, default=0) parser.add_argument('--save_path', type=str, default='./save') parser.add_argument('--save_file_name', type=str, default='vgg_cifar') return parser def main(opts): # 1. set random seeds set_random_seeds(random_seed=0) # 2. initialization init_for_distributed(opts) # 3. visdom vis = None # 4. data set transform_train = tfs.Compose([ tfs.Resize(256), tfs.RandomCrop(224), tfs.RandomHorizontalFlip(), tfs.ToTensor(), tfs.Normalize(mean=(0.4914, 0.4822, 0.4465), std=(0.2023, 0.1994, 0.2010)), ]) transform_test = tfs.Compose([ tfs.Resize(256), tfs.CenterCrop(224), tfs.ToTensor(), tfs.Normalize(mean=(0.4914, 0.4822, 0.4465), std=(0.2023, 0.1994, 0.2010)), ]) train_set = CIFAR10(root=opts.root, train=True, transform=transform_train, download=True) test_set = CIFAR10(root=opts.root, train=False, transform=transform_test, download=True) train_sampler = DistributedSampler(dataset=train_set, shuffle=True) test_sampler = DistributedSampler(dataset=test_set, shuffle=False) train_loader = DataLoader(dataset=train_set, batch_size=int(opts.batch_size / opts.world_size), shuffle=False, num_workers=int(opts.num_workers / opts.world_size), sampler=train_sampler, pin_memory=True) test_loader = DataLoader(dataset=test_set, batch_size=int(opts.batch_size / opts.world_size), shuffle=False, num_workers=int(opts.num_workers / opts.world_size), sampler=test_sampler, pin_memory=True) # 5. model model = vgg11(pretrained=False) model = model.cuda(opts.local_rank) model = DDP(module=model, device_ids=[opts.local_rank]) # 6. criterion criterion = torch.nn.CrossEntropyLoss().to(opts.local_rank) # 7. optimizer optimizer = torch.optim.SGD(params=model.parameters(), lr=0.01, weight_decay=0.0005, momentum=0.9) # 8. scheduler scheduler = StepLR(optimizer=optimizer, step_size=30, gamma=0.1) if opts.start_epoch != 0: checkpoint = torch.load(os.path.join(opts.save_path, opts.save_file_name) + '.{}.pth.tar' .format(opts.start_epoch - 1), map_location=torch.device('cuda:{}'.format(opts.local_rank))) model.load_state_dict(checkpoint['model_state_dict']) optimizer.load_state_dict(checkpoint['optimizer_state_dict']) scheduler.load_state_dict(checkpoint['scheduler_state_dict']) if opts.global_rank == 0: print('\nLoaded checkpoint from epoch %d.\n' % (int(opts.start_epoch) - 1)) for epoch in range(opts.start_epoch, opts.epoch): # 9. train tic = time.time() model.train() train_sampler.set_epoch(epoch) for i, (images, labels) in enumerate(train_loader): images = images.to(opts.local_rank) labels = labels.to(opts.local_rank) outputs = model(images) # ----------- update ----------- optimizer.zero_grad() loss = criterion(outputs, labels) loss.backward() optimizer.step() # get lr for param_group in optimizer.param_groups: lr = param_group['lr'] # time toc = time.time() # visualization if (i % opts.vis_step == 0 or i == len(train_loader) - 1): print('GPU[{0}] Epoch [{1}/{2}], Iter [{3}/{4}], Loss: {5:.4f}, LR: {6:.5f}, Time: {7:.2f}'.format(opts.global_rank, epoch, opts.epoch, i, len(train_loader), loss.item(), lr, toc - tic)) # save pth file if opts.local_rank == 0: if not os.path.exists(opts.save_path): os.mkdir(opts.save_path) checkpoint = {'epoch': epoch, 'model_state_dict': model.state_dict(), 'optimizer_state_dict': optimizer.state_dict(), 'scheduler_state_dict': scheduler.state_dict()} torch.save(checkpoint, os.path.join(opts.save_path, opts.save_file_name + '.{}.pth.tar'.format(epoch))) print("save pth.tar {} epoch!".format(epoch)) # 10. test model.eval() val_avg_loss = 0 correct_top1 = 0 correct_top5 = 0 total = 0 with torch.no_grad(): for i, (images, labels) in enumerate(test_loader): images = images.to(opts.local_rank) labels = labels.to(opts.local_rank) outputs = model(images) loss = criterion(outputs, labels) val_avg_loss += loss.item() # rank 1 _, pred = torch.max(outputs, 1) total += labels.size(0) correct_top1 += (pred == labels).sum().item() # rank 5 _, rank5 = outputs.topk(5, 1, True, True) rank5 = rank5.t() correct5 = rank5.eq(labels.view(1, -1).expand_as(rank5)) for k in range(5): correct_k = correct5[:k+1].reshape(-1).float().sum(0, keepdim=True) correct_top5 += correct_k.item() accuracy_top1 = correct_top1 / total accuracy_top5 = correct_top5 / total val_avg_loss = val_avg_loss / len(test_loader) print("top-1 percentage : {0:0.3f}%".format(correct_top1 / total * 100)) print("top-5 percentage : {0:0.3f}%".format(correct_top5 / total * 100)) scheduler.step() return 0 def init_for_distributed(opts): # 1. setting for distributed training opts.global_rank = int(os.environ['RANK']) opts.local_rank = int(os.environ['LOCAL_RANK']) opts.world_size = int(os.environ['WORLD_SIZE']) torch.cuda.set_device(opts.local_rank) if opts.global_rank is not None and opts.local_rank is not None: print("Use GPU: [{}/{}] for training".format(opts.global_rank, opts.local_rank)) # 2. init_process_group dist.init_process_group( backend="nccl", rank=opts.global_rank, world_size=opts.world_size, device_id=torch.device(f"cuda:{opts.local_rank}"), timeout=datetime.timedelta(seconds=60) ) return if __name__ == '__main__': parser = argparse.ArgumentParser('vgg11 cifar training', parents=[get_args_parser()]) opts = parser.parse_args() main(opts) --- # 1. Volcano Queue for training workloads # Queue는 리소스 할당 단위. capability는 클러스터 실제 용량에 맞게 조정 필요. apiVersion: scheduling.volcano.sh/v1beta1 kind: Queue metadata: name: training-queue spec: weight: 1 reclaimable: true capability: # 3노드 x GPU 8장 = 24, hostdev 2개 = 6 cpu: "100" memory: "512Gi" nvidia.com/gpu: "24" nvidia.com/hostdev: "6" --- # 2. ClusterTrainingRuntime with Volcano gang scheduling + topology-aware scheduling # podGroupPolicy.volcano를 사용하면 PodGroup이 자동 생성됨 (수동 생성 불필요) # MASTER_ADDR / MASTER_PORT는 Kubeflow Trainer torch runtime이 자동 설정 (torchrun rdzv) apiVersion: trainer.kubeflow.org/v1alpha1 kind: ClusterTrainingRuntime metadata: name: torch-distributed-volcano labels: trainer.kubeflow.org/framework: torch spec: mlPolicy: torch: numProcPerNode: 8 numNodes: 2 # Volcano gang scheduling 활성화 - PodGroup 자동 생성 podGroupPolicy: volcano: networkTopology: mode: hard highestTierAllowed: 1 template: metadata: annotations: # Queue 지정 (runtime level) scheduling.volcano.sh/queue-name: training-queue spec: replicatedJobs: - name: node template: metadata: labels: trainer.kubeflow.org/trainjob-ancestor-step: trainer spec: template: metadata: annotations: # InfiniBand CNI k8s.v1.cni.cncf.io/networks: hostdevice-net spec: containers: - name: node image: nvcr.io/nvidia/pytorch:24.10-py3 securityContext: privileged: false capabilities: add: - IPC_LOCK env: - name: NCCL_DEBUG value: "INFO" - name: NCCL_DEBUG_SUBSYS value: "INIT,NET,IB" - name: NCCL_IB_DISABLE value: "0" - name: NCCL_SOCKET_IFNAME value: "net" resources: requests: nvidia.com/gpu: "8" nvidia.com/hostdev: "2" limits: nvidia.com/gpu: "8" nvidia.com/hostdev: "2" volumeMounts: - name: shared-memory mountPath: /dev/shm - name: cifar-data mountPath: /workspace/data/cifar-10-batches-py - name: training-scripts mountPath: /workspace/scripts volumes: - name: shared-memory emptyDir: medium: Memory sizeLimit: 128Gi - name: cifar-data hostPath: path: /home/ubuntu/cifar-10-batches-py type: Directory - name: training-scripts configMap: name: training-scripts defaultMode: 0755 --- # 3. Example TrainJob (테스트용 - 2노드 분산학습) # podGroupPolicy, env, volumes, command 등은 runtime에서 상속됨 apiVersion: trainer.kubeflow.org/v1alpha1 kind: TrainJob metadata: name: example-distributed-training namespace: default spec: runtimeRef: name: torch-distributed-volcano trainer: image: nvcr.io/nvidia/pytorch:24.10-py3 command: - torchrun - /workspace/scripts/train_nccl.py numNodes: 2 resourcesPerNode: requests: nvidia.com/gpu: "8" nvidia.com/hostdev: "2" limits: nvidia.com/gpu: "8" nvidia.com/hostdev: "2"