From 63fbcf8085bada628558b4d60acfae0066288e26 Mon Sep 17 00:00:00 2001 From: Cloud User Date: Tue, 3 Mar 2026 10:24:13 +0900 Subject: [PATCH] Add TensorFlow custom runtime for Kubeflow Trainer v2 with Volcano integration MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Trainer v2에는 TensorFlow runtime이 기본 제공되지 않으므로 Custom ClusterTrainingRuntime을 생성하여 MultiWorkerMirroredStrategy 기반 분산학습을 지원한다. Pod hostname + Headless Service DNS로 TF_CONFIG를 자동 구성. Co-Authored-By: Claude Opus 4.6 --- .gitignore | 1 + README.md | 211 ++++++++++ tensorflow-custom-runtime.yaml | 46 +++ tensorflow-volcano-trainjob-integration.yaml | 318 +++++++++++++++ volcano-trainjob-integration.yaml | 392 +++++++++++++++++++ 5 files changed, 968 insertions(+) create mode 100644 .gitignore create mode 100644 README.md create mode 100644 tensorflow-custom-runtime.yaml create mode 100644 tensorflow-volcano-trainjob-integration.yaml create mode 100644 volcano-trainjob-integration.yaml diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..4c49bd7 --- /dev/null +++ b/.gitignore @@ -0,0 +1 @@ +.env diff --git a/README.md b/README.md new file mode 100644 index 0000000..21f8ef9 --- /dev/null +++ b/README.md @@ -0,0 +1,211 @@ +# Kubeflow Trainer v2 - TensorFlow Custom Runtime + +Kubeflow Trainer v2에서 TensorFlow 분산학습을 위한 Custom ClusterTrainingRuntime 생성 및 Volcano 연동 가이드. + +## 배경 + +Kubeflow Trainer v2 (v1alpha1)에는 다음 런타임만 기본 제공된다: + +| Runtime | Framework Label | mlPolicy 필드 | 엔트리포인트 주입 | +|---------|----------------|---------------|-------------------| +| torch-distributed | `torch` | `mlPolicy.torch` | torchrun + MASTER_ADDR/RANK 등 자동 설정 | +| deepspeed-distributed | `deepspeed` | `mlPolicy.mpi` | MPI (OpenMPI) + SSH 자동 설정 | +| mlx-distributed | `mlx` | `mlPolicy.mpi` | MPI (OpenMPI) + SSH 자동 설정 | +| torchtune-* | `torchtune` | `mlPolicy.torch` | tune run + rdzv 자동 설정 | + +**TensorFlow runtime은 없다.** 따라서 Custom ClusterTrainingRuntime을 직접 생성해야 한다. + +## PyTorch vs TensorFlow Runtime 핵심 차이 + +``` +┌─────────────────────────────────────────────────────────────────────┐ +│ PyTorch Runtime (built-in) │ +│ │ +│ mlPolicy.torch → 컨트롤러가 자동으로: │ +│ ├── torchrun 엔트리포인트 주입 │ +│ ├── MASTER_ADDR, MASTER_PORT 설정 │ +│ ├── RANK, LOCAL_RANK, WORLD_SIZE 설정 │ +│ └── rdzv (rendezvous) 자동 구성 │ +│ │ +│ → 학습 코드에서 분산환경 설정 불필요 │ +└─────────────────────────────────────────────────────────────────────┘ + +┌─────────────────────────────────────────────────────────────────────┐ +│ TensorFlow Runtime (custom) │ +│ │ +│ mlPolicy에 framework 없음 → 컨트롤러가 주입하는 것 없음 │ +│ │ +│ 학습 스크립트에서 자체 구성: │ +│ ├── Pod hostname 파싱 → TrainJob 이름, worker index 추출 │ +│ ├── Headless Service DNS로 worker 주소 목록 생성 │ +│ ├── TF_CONFIG 환경변수 자동 구성 │ +│ └── tf.distribute.MultiWorkerMirroredStrategy 사용 │ +│ │ +│ 필수 설정: │ +│ ├── network.publishNotReadyAddresses: true (DNS 조기 해석) │ +│ └── NUM_WORKERS 환경변수 = mlPolicy.numNodes │ +└─────────────────────────────────────────────────────────────────────┘ +``` + +## TF_CONFIG 자동 구성 원리 + +Trainer v2에서 TrainJob을 생성하면 내부적으로 JobSet이 만들어지고, 각 Pod은 예측 가능한 hostname을 갖는다: + +``` +Pod hostname 패턴: {trainjob-name}-node-{replica_index}-{completion_index} + +예시 (TrainJob: tf-distributed-training, numNodes: 2): + Worker 0: tf-distributed-training-node-0-0 + Worker 1: tf-distributed-training-node-1-0 +``` + +Headless Service DNS와 결합하면: + +``` +Worker 0: tf-distributed-training-node-0-0.tf-distributed-training.default.svc.cluster.local:12345 +Worker 1: tf-distributed-training-node-1-0.tf-distributed-training.default.svc.cluster.local:12345 +``` + +학습 스크립트가 이를 파싱하여 자동으로 TF_CONFIG를 생성한다: + +```json +{ + "cluster": { + "worker": [ + "tf-distributed-training-node-0-0.tf-distributed-training.default.svc.cluster.local:12345", + "tf-distributed-training-node-1-0.tf-distributed-training.default.svc.cluster.local:12345" + ] + }, + "task": {"type": "worker", "index": 0} +} +``` + +## 파일 구조 + +``` +. +├── README.md # 이 문서 +├── tensorflow-custom-runtime.yaml # TF Custom Runtime (기본, Volcano 미포함) +├── tensorflow-volcano-trainjob-integration.yaml # TF + Volcano 전체 통합 (ConfigMap + Queue + Runtime + TrainJob) +└── volcano-trainjob-integration.yaml # (참고) PyTorch + Volcano 통합 원본 +``` + +## 적용 방법 + +### 1단계: 기본 TensorFlow Runtime만 적용 + +Volcano 없이 기본 TF 런타임만 필요한 경우: + +```bash +kubectl apply -f tensorflow-custom-runtime.yaml +``` + +확인: +```bash +kubectl get clustertrainingruntime +# tensorflow-distributed 가 목록에 나타나야 함 +``` + +### 2단계: Volcano 연동 전체 적용 + +Volcano gang scheduling + 토폴로지 인식 스케줄링이 포함된 전체 버전: + +```bash +# ConfigMap (학습 스크립트) + Queue + Runtime + TrainJob 한번에 적용 +kubectl apply -f tensorflow-volcano-trainjob-integration.yaml +``` + +확인: +```bash +# Runtime 확인 +kubectl get clustertrainingruntime tensorflow-distributed-volcano + +# TrainJob 상태 확인 +kubectl get trainjob tf-distributed-training + +# Pod 상태 확인 +kubectl get pods -l batch.kubernetes.io/job-name + +# Volcano PodGroup 확인 +kubectl get podgroup + +# 로그 확인 (worker 0) +kubectl logs -l batch.kubernetes.io/job-name=tf-distributed-training-node-0 -f +``` + +### 3단계: TrainJob만 변경하여 재실행 + +Runtime은 유지하고 TrainJob만 변경할 경우: + +```bash +# 기존 TrainJob 삭제 +kubectl delete trainjob tf-distributed-training + +# 수정 후 재적용 +kubectl apply -f tensorflow-volcano-trainjob-integration.yaml +``` + +## 주의사항 + +### NUM_WORKERS와 numNodes 동기화 + +`NUM_WORKERS` 환경변수는 `mlPolicy.numNodes` (또는 TrainJob의 `trainer.numNodes`)와 **반드시 일치**해야 한다. + +PyTorch runtime은 컨트롤러가 `WORLD_SIZE`를 자동 주입하지만, 커스텀 TF runtime은 이를 수동으로 관리해야 한다. + +```yaml +# Runtime에서 기본값 설정 +spec: + mlPolicy: + numNodes: 2 # ← 이 값과 + ... + env: + - name: NUM_WORKERS + value: "2" # ← 이 값이 일치해야 함 +``` + +TrainJob에서 numNodes를 오버라이드할 경우, env도 함께 오버라이드해야 한다: + +```yaml +apiVersion: trainer.kubeflow.org/v1alpha1 +kind: TrainJob +spec: + runtimeRef: + name: tensorflow-distributed-volcano + trainer: + numNodes: 4 # 4노드로 변경 + env: + - name: NUM_WORKERS + value: "4" # 반드시 함께 변경 +``` + +### InfiniBand 설정 + +클러스터에 InfiniBand가 없는 경우 다음 항목을 제거해야 한다: + +```yaml +# 제거 대상: +metadata: + annotations: + k8s.v1.cni.cncf.io/networks: hostdevice-net # ← 제거 + +resources: + requests: + nvidia.com/hostdev: "2" # ← 제거 + limits: + nvidia.com/hostdev: "2" # ← 제거 +``` + +### network.publishNotReadyAddresses + +TensorFlow의 `MultiWorkerMirroredStrategy`는 시작 시 모든 worker에 gRPC 연결을 시도한다. 모든 Pod이 Ready 상태가 되기 전에도 DNS가 해석되어야 하므로 이 설정이 **필수**다: + +```yaml +spec: + template: + spec: + network: + publishNotReadyAddresses: true # 반드시 true +``` + +PyTorch runtime에서는 torchrun의 rdzv(rendezvous) 메커니즘이 이를 처리하므로 불필요하다. diff --git a/tensorflow-custom-runtime.yaml b/tensorflow-custom-runtime.yaml new file mode 100644 index 0000000..a4f6836 --- /dev/null +++ b/tensorflow-custom-runtime.yaml @@ -0,0 +1,46 @@ +# Kubeflow Trainer v2 - TensorFlow Custom ClusterTrainingRuntime +# Trainer v2에는 TensorFlow runtime이 기본 제공되지 않으므로 Custom Runtime으로 생성 +# +# 핵심 차이점: +# - PyTorch runtime: mlPolicy.torch 사용 → torchrun 자동 주입, MASTER_ADDR/RANK 등 자동 설정 +# - TensorFlow runtime (custom): mlPolicy에 framework 없음 → TF_CONFIG를 학습 스크립트에서 자체 구성 +# +# TF_CONFIG 자동 구성 방식: +# 1. Pod hostname에서 TrainJob 이름과 replica index 추출 +# 2. Headless Service DNS를 이용한 worker 주소 목록 생성 +# 3. tf.distribute.MultiWorkerMirroredStrategy로 분산학습 실행 + +apiVersion: trainer.kubeflow.org/v1alpha1 +kind: ClusterTrainingRuntime +metadata: + name: tensorflow-distributed + labels: + trainer.kubeflow.org/framework: tensorflow +spec: + mlPolicy: + numNodes: 1 + template: + spec: + network: + publishNotReadyAddresses: true + replicatedJobs: + - name: node + template: + metadata: + labels: + trainer.kubeflow.org/trainjob-ancestor-step: trainer + spec: + template: + spec: + containers: + - name: node + image: tensorflow/tensorflow:2.16.1-gpu + env: + - name: NUM_WORKERS + value: "1" + - name: TF_WORKER_PORT + value: "12345" + - name: NAMESPACE + valueFrom: + fieldRef: + fieldPath: metadata.namespace diff --git a/tensorflow-volcano-trainjob-integration.yaml b/tensorflow-volcano-trainjob-integration.yaml new file mode 100644 index 0000000..7ae9f2f --- /dev/null +++ b/tensorflow-volcano-trainjob-integration.yaml @@ -0,0 +1,318 @@ +# Volcano <-> Kubeflow Trainer v2 TensorFlow Integration +# Reference: https://www.kubeflow.org/docs/components/trainer/gang-scheduling/volcano/ +# +# Trainer v2에는 TensorFlow runtime이 없으므로 Custom Runtime을 생성하여 사용 +# TF_CONFIG는 학습 스크립트에서 Pod hostname + Headless Service DNS로 자동 구성 + +--- +# 0. ConfigMap for TensorFlow training script (CNN + CIFAR-10 + MultiWorkerMirroredStrategy) +apiVersion: v1 +kind: ConfigMap +metadata: + name: tf-training-scripts +data: + train_tf.py: | + """ + TensorFlow Distributed Training Example for Kubeflow Trainer v2 + - CNN (CIFAR-10) + MultiWorkerMirroredStrategy + NCCL + - TF_CONFIG를 Pod hostname과 환경변수로 자동 구성 + """ + import os + import json + import socket + import re + import tensorflow as tf + import numpy as np + + + def setup_tf_config(): + """ + Kubeflow Trainer v2의 Pod 네이밍 규칙을 이용하여 TF_CONFIG를 자동 구성. + + Pod hostname 패턴: {trainjob-name}-node-{replica_index}-{completion_index} + Headless Service DNS: {hostname}.{trainjob-name}.{namespace}.svc.cluster.local + + 환경변수: + - NUM_WORKERS: 총 워커 수 (runtime의 numNodes와 일치해야 함) + - TF_WORKER_PORT: gRPC 통신 포트 (기본값: 12345) + - NAMESPACE: Pod이 실행되는 namespace (Downward API) + """ + hostname = socket.gethostname() + num_workers = int(os.environ.get('NUM_WORKERS', '1')) + port = os.environ.get('TF_WORKER_PORT', '12345') + namespace = os.environ.get('NAMESPACE', 'default') + + # hostname에서 TrainJob 이름과 replica index 추출 + # 패턴: {trainjob-name}-node-{replica_index}-{completion_index} + match = re.match(r'^(.+)-node-(\d+)-(\d+)$', hostname) + if match: + job_name = match.group(1) + worker_index = int(match.group(2)) + else: + print(f"[WARNING] Cannot parse hostname '{hostname}', using defaults") + job_name = os.environ.get('TRAINJOB_NAME', 'unknown') + worker_index = 0 + + # Headless Service DNS를 이용한 worker 주소 목록 생성 + workers = [] + for i in range(num_workers): + worker_host = ( + f"{job_name}-node-{i}-0.{job_name}.{namespace}" + f".svc.cluster.local:{port}" + ) + workers.append(worker_host) + + tf_config = { + "cluster": {"worker": workers}, + "task": {"type": "worker", "index": worker_index} + } + + os.environ['TF_CONFIG'] = json.dumps(tf_config) + print(f"[Worker {worker_index}/{num_workers}] hostname: {hostname}") + print(f"[Worker {worker_index}/{num_workers}] TF_CONFIG:") + print(json.dumps(tf_config, indent=2)) + + return worker_index, num_workers + + + def build_model(): + """CIFAR-10용 CNN 모델""" + model = tf.keras.Sequential([ + tf.keras.layers.Conv2D( + 32, (3, 3), activation='relu', input_shape=(32, 32, 3)), + tf.keras.layers.BatchNormalization(), + tf.keras.layers.Conv2D(32, (3, 3), activation='relu'), + tf.keras.layers.MaxPooling2D((2, 2)), + tf.keras.layers.Dropout(0.25), + + tf.keras.layers.Conv2D(64, (3, 3), activation='relu'), + tf.keras.layers.BatchNormalization(), + tf.keras.layers.Conv2D(64, (3, 3), activation='relu'), + tf.keras.layers.MaxPooling2D((2, 2)), + tf.keras.layers.Dropout(0.25), + + tf.keras.layers.Flatten(), + tf.keras.layers.Dense(512, activation='relu'), + tf.keras.layers.BatchNormalization(), + tf.keras.layers.Dropout(0.5), + tf.keras.layers.Dense(10) + ]) + return model + + + def main(): + worker_index, num_workers = setup_tf_config() + + # MultiWorkerMirroredStrategy (GPU간 NCCL 통신) + communication_options = tf.distribute.experimental.CommunicationOptions( + implementation=tf.distribute.experimental.CommunicationImplementation.NCCL + ) + strategy = tf.distribute.MultiWorkerMirroredStrategy( + communication_options=communication_options + ) + + print(f"[Worker {worker_index}] num_replicas_in_sync: " + f"{strategy.num_replicas_in_sync}") + + # CIFAR-10 데이터셋 로드 + (x_train, y_train), (x_test, y_test) = ( + tf.keras.datasets.cifar10.load_data() + ) + x_train = x_train.astype('float32') / 255.0 + x_test = x_test.astype('float32') / 255.0 + + # 배치 크기 설정 + BATCH_SIZE_PER_REPLICA = 64 + GLOBAL_BATCH_SIZE = ( + BATCH_SIZE_PER_REPLICA * strategy.num_replicas_in_sync + ) + + # tf.data.Dataset 생성 + train_dataset = tf.data.Dataset.from_tensor_slices( + (x_train, y_train) + ).shuffle(50000).batch(GLOBAL_BATCH_SIZE).prefetch(tf.data.AUTOTUNE) + + test_dataset = tf.data.Dataset.from_tensor_slices( + (x_test, y_test) + ).batch(GLOBAL_BATCH_SIZE).prefetch(tf.data.AUTOTUNE) + + # Strategy scope 내에서 모델 생성 및 컴파일 + with strategy.scope(): + model = build_model() + model.compile( + optimizer=tf.keras.optimizers.Adam(learning_rate=0.001), + loss=tf.keras.losses.SparseCategoricalCrossentropy( + from_logits=True), + metrics=['accuracy'] + ) + + # 학습 + EPOCHS = int(os.environ.get('EPOCHS', '10')) + verbose = 2 if worker_index == 0 else 0 + + callbacks = [] + if worker_index == 0: + callbacks.append( + tf.keras.callbacks.TensorBoard(log_dir='/workspace/logs') + ) + + model.fit( + train_dataset, + epochs=EPOCHS, + validation_data=test_dataset, + verbose=verbose, + callbacks=callbacks + ) + + # 평가 (worker 0만 출력) + if worker_index == 0: + loss, accuracy = model.evaluate(test_dataset, verbose=2) + print(f"\n===== Final Results =====") + print(f"Test Loss: {loss:.4f}") + print(f"Test Accuracy: {accuracy:.4f}") + print(f"=========================") + + + if __name__ == '__main__': + main() + +--- +# 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: + # 클러스터 리소스에 맞게 조정 + cpu: "100" + memory: "512Gi" + nvidia.com/gpu: "24" + nvidia.com/hostdev: "6" + +--- +# 2. TensorFlow Custom ClusterTrainingRuntime with Volcano gang scheduling +# +# PyTorch runtime과의 핵심 차이: +# - mlPolicy에 torch/mpi 필드 없음 → 컨트롤러가 torchrun/MPI 주입하지 않음 +# - network.publishNotReadyAddresses: true → Pod간 DNS 조기 해석 가능 +# - TF_CONFIG는 학습 스크립트에서 hostname + Headless Service DNS로 자동 구성 +# - NUM_WORKERS 환경변수가 numNodes와 반드시 일치해야 함 +apiVersion: trainer.kubeflow.org/v1alpha1 +kind: ClusterTrainingRuntime +metadata: + name: tensorflow-distributed-volcano + labels: + trainer.kubeflow.org/framework: tensorflow +spec: + mlPolicy: + 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: + network: + # TensorFlow MultiWorkerMirroredStrategy가 gRPC로 통신하므로 + # 모든 Pod이 Ready 전에도 DNS 해석 가능해야 함 + publishNotReadyAddresses: true + replicatedJobs: + - name: node + template: + metadata: + labels: + trainer.kubeflow.org/trainjob-ancestor-step: trainer + spec: + template: + metadata: + annotations: + # InfiniBand CNI (클러스터에 InfiniBand가 없으면 이 줄 제거) + k8s.v1.cni.cncf.io/networks: hostdevice-net + spec: + containers: + - name: node + image: nvcr.io/nvidia/tensorflow:24.03-tf2-py3 + securityContext: + privileged: true + capabilities: + add: + - IPC_LOCK + env: + # --- NCCL 설정 (TF MultiWorkerMirroredStrategy + NCCL backend) --- + - 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" + # --- TF_CONFIG 자동 구성용 환경변수 --- + # NUM_WORKERS는 mlPolicy.numNodes와 반드시 일치해야 함 + - name: NUM_WORKERS + value: "2" + - name: TF_WORKER_PORT + value: "12345" + - name: NAMESPACE + valueFrom: + fieldRef: + fieldPath: metadata.namespace + # 학습 에포크 수 + - name: EPOCHS + value: "10" + 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: training-scripts + mountPath: /workspace/scripts + volumes: + - name: shared-memory + emptyDir: + medium: Memory + sizeLimit: 128Gi + - name: training-scripts + configMap: + name: tf-training-scripts + defaultMode: 0755 + +--- +# 3. Example TrainJob (TensorFlow 2노드 분산학습) +# podGroupPolicy, env, volumes, command 등은 runtime에서 상속됨 +apiVersion: trainer.kubeflow.org/v1alpha1 +kind: TrainJob +metadata: + name: tf-distributed-training + namespace: default +spec: + runtimeRef: + name: tensorflow-distributed-volcano + trainer: + image: nvcr.io/nvidia/tensorflow:24.03-tf2-py3 + command: + - python + - /workspace/scripts/train_tf.py + numNodes: 2 + resourcesPerNode: + requests: + nvidia.com/gpu: "8" + nvidia.com/hostdev: "2" + limits: + nvidia.com/gpu: "8" + nvidia.com/hostdev: "2" diff --git a/volcano-trainjob-integration.yaml b/volcano-trainjob-integration.yaml new file mode 100644 index 0000000..d22e483 --- /dev/null +++ b/volcano-trainjob-integration.yaml @@ -0,0 +1,392 @@ +# 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: true + 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" +