Fine-Tuning LoRA Adapters on Kubernetes: Orchestrating Distributed Training with Ray and Ludwig — code snippets

snippet-1.yaml

apiVersion: ray.io/v1
kind: RayCluster
metadata:
  name: ray-cluster-lora
  namespace: ml-platform
spec:
  rayVersion: '2.35.0'
  headGroupSpec:
    rayStartParams:
      dashboard-host: '0.0.0.0'
    template:
      spec:
        containers:
        - name: ray-head
          image: rayproject/ray:2.35.0-py310
          resources:
            limits:
              cpu: "8"
              memory: "32Gi"
            requests:
              cpu: "8"
              memory: "32Gi"
          volumeMounts:
          - mountPath: /dev/shm
            name: dshm
        volumes:
        - name: dshm
          emptyDir:
            medium: Memory
            sizeLimit: 16Gi
  workerGroupSpecs:
  - groupName: gpu-workers
    replicas: 2
    minReplicas: 1
    maxReplicas: 4
    rayStartParams: {}
    template:
      spec:
        containers:
        - name: ray-worker
          image: rayproject/ray:2.35.0-py310-gpu
          resources:
            limits:
              cpu: "16"
              memory: "64Gi"
              nvidia.com/gpu: "2"
            requests:
              cpu: "16"
              memory: "64Gi"
              nvidia.com/gpu: "2"
          volumeMounts:
          - mountPath: /dev/shm
            name: dshm
        volumes:
        - name: dshm
          emptyDir:
            medium: Memory
            sizeLimit: 32Gi
        tolerations:
        - key: "nvidia.com/gpu"
          operator: "Exists"
          effect: "NoSchedule"
        nodeSelector:
          node.kubernetes.io/instance-type: "g5.4xlarge"

snippet-2.yaml

model_type: llm
base_model: meta-llama/Meta-Llama-3-8B-Instruct
input_features:
  - name: instruction
    type: text
output_features:
  - name: response
    type: text
prompt:
  template: |
    Below is an instruction that describes a task, paired with an input that provides further context.
    Write a response that appropriately completes the request.

    ### Instruction:
    {instruction}

    ### Response:
    {response}
adapter:
  type: lora
  r: 8
  alpha: 16
  dropout: 0.05
  target_modules:
    - q_proj
    - v_proj
    - k_proj
    - o_proj
trainer:
  type: llm
  epochs: 3
  batch_size: 2
  eval_batch_size: 2
  gradient_accumulation_steps: 8
  learning_rate: 0.0002
  optimizer:
    type: adamw
  lr_scheduler:
    name: cosine
backend:
  type: ray
  cache_dir: /tmp/ludwig_cache
  trainer:
    use_gpu: true
    num_workers: 4
    resources_per_worker:
      CPU: 4
      GPU: 1

snippet-3.py

import logging
import ray
from ludwig.api import LudwigModel

logging.basicConfig(level=logging.INFO)

# Initialize connection to the active Ray cluster context
ray.init(address="auto", ignore_reinit_error=True)

# Define dataset paths in cloud object storage
training_set = "s3://ml-ops-data/train.parquet"
validation_set = "s3://ml-ops-data/validation.parquet"

# Initialize Ludwig model using the configuration file
config_path = "ludwig_config.yaml"
model = LudwigModel(config=config_path, logging_level=logging.INFO)

# Run the training pipeline across the Ray cluster nodes
results = model.train(
    dataset=training_set,
    validation_dataset=validation_set,
    output_directory="s3://ml-ops-models/lora-finetuning-results"
)

print("Training finished successfully. Metrics:")
print(results.train_stats)

snippet-4.py

import time
from ray.job_submission import JobSubmissionClient, JobStatus

# Connect to the Ray Head Node dashboard service via internal DNS or port-forwarding
client = JobSubmissionClient("http://ray-cluster-head-svc.ml-platform.svc.cluster.local:8265")

# Define job requirements and environment variables
job_id = client.submit_job(
    entrypoint="python train.py",
    runtime_env={
        "working_dir": "./",
        "pip": [
            "ludwig[llm]==0.8.6",
            "deepspeed==0.12.3",
            "pyarrow==14.0.1",
            "fsspec[s3]==2023.10.0"
        ],
        "env_vars": {
            "NCCL_DEBUG": "INFO",
            "NCCL_IB_DISABLE": "1",
            "AWS_DEFAULT_REGION": "us-west-2"
        }
    }
)

print(f"Job submitted. Job ID: {job_id}")

# Poll the job status until it reaches a terminal state
while True:
    status_info = client.get_job_status(job_id)
    print(f"Job status: {status_info.status}")
    if status_info.status in [JobStatus.SUCCEEDED, JobStatus.FAILED, JobStatus.STOPPED]:
        break
    time.sleep(10)

snippet-5.json

{
  "fp16": {
    "enabled": "auto"
  },
  "bf16": {
    "enabled": "auto"
  },
  "zero_optimization": {
    "stage": 3,
    "offload_optimizer": {
      "device": "cpu",
      "pin_memory": true
    },
    "offload_param": {
      "device": "cpu",
      "pin_memory": true
    },
    "overlap_comm": true,
    "contiguous_gradients": true,
    "reduce_bucket_size": "auto",
    "stage3_prefetch_bucket_size": "auto",
    "stage3_param_persistence_threshold": "auto"
  },
  "gradient_accumulation_steps": "auto",
  "gradient_clipping": "auto",
  "train_batch_size": "auto",
  "train_micro_batch_size_per_gpu": "auto"
}

snippet-6.yaml

apiVersion: apps/v1
kind: Deployment
metadata:
  name: vllm-lora-service
  namespace: ml-platform
spec:
  replicas: 2
  selector:
    matchLabels:
      app: vllm-lora-service
  template:
    metadata:
      labels:
        app: vllm-lora-service
    spec:
      containers:
      - name: vllm-container
        image: vllm/vllm-openai:v0.4.2
        command: ["python3", "-m", "vllm.entrypoints.openai.api_server"]
        args:
        - "--model"
        - "meta-llama/Meta-Llama-3-8B-Instruct"
        - "--enable-lora"
        - "--lora-modules"
        - "customer-support-adapter=s3://ml-ops-models/lora-finetuning-results/model/model_weights"
        - "--max-lora-rank"
        - "8"
        - "--port"
        - "8000"
        ports:
        - containerPort: 8000
        resources:
          limits:
            cpu: "8"
            memory: "32Gi"
            nvidia.com/gpu: "1"
          requests:
            cpu: "8"
            memory: "32Gi"
            nvidia.com/gpu: "1"