GPU & ML Training Across Clouds

# Launch a SageMaker training job with multiple GPUs
aws sagemaker create-training-job \
  --training-job-name "llm-finetune-$(date +%Y%m%d)" \
  --algorithm-specification '{
    "TrainingImage": "763104351884.dkr.ecr.us-east-1.amazonaws.com/pytorch-training:2.1.0-gpu-py310-cu121-ubuntu20.04-sagemaker",
    "TrainingInputMode": "FastFile"
  }' \
  --resource-config '{
    "InstanceType": "ml.p4de.24xlarge",
    "InstanceCount": 4,
    "VolumeSizeInGB": 500
  }' \
  --input-data-config '[{
    "ChannelName": "training",
    "DataSource": {"S3DataSource": {"S3Uri": "s3://training-data/dataset/", "S3DataType": "S3Prefix"}},
    "InputMode": "FastFile"
  }]' \
  --output-data-config '{"S3OutputPath": "s3://model-artifacts/"}' \
  --role-arn "arn:aws:iam::123456789012:role/SageMakerRole" \
  --stopping-condition '{"MaxRuntimeInSeconds": 86400}'

# Or launch EC2 GPU instances directly
aws ec2 run-instances \
  --instance-type p4de.24xlarge \
  --image-id ami-deep-learning \
  --key-name ml-keypair \
  --count 4 \
  --placement '{"GroupName": "ml-cluster", "Tenancy": "default"}' \
  --network-interfaces '[{"DeviceIndex": 0, "SubnetId": "subnet-abc123", "Groups": ["sg-ml"]}]'

# Submit a Vertex AI custom training job
gcloud ai custom-jobs create \
  --display-name="llm-finetune" \
  --region=us-central1 \
  --worker-pool-spec=machine-type=a2-ultragpu-8g,replica-count=4,accelerator-type=NVIDIA_A100_80GB,accelerator-count=8,container-image-uri=us-docker.pkg.dev/PROJECT/repo/training:latest

# Create a GKE node pool with GPUs
gcloud container node-pools create gpu-pool \
  --cluster=ml-cluster \
  --zone=us-central1-a \
  --machine-type=a2-ultragpu-8g \
  --accelerator=type=nvidia-a100-80gb,count=8 \
  --num-nodes=4 \
  --min-nodes=0 \
  --max-nodes=8 \
  --enable-autoscaling \
  --spot

# Create an Azure ML compute cluster
az ml compute create \
  --name gpu-cluster \
  --type AmlCompute \
  --size Standard_ND96amsr_A100_v4 \
  --min-instances 0 \
  --max-instances 8 \
  --idle-time-before-scale-down 600 \
  --resource-group ml-rg \
  --workspace-name ml-workspace

# Submit a training job
az ml job create --file training-job.yml \
  --resource-group ml-rg \
  --workspace-name ml-workspace

# PyTorch distributed training with DeepSpeed
# launch command: deepspeed --num_gpus=8 --num_nodes=4 train.py

import deepspeed
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

def train():
    model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-3-8B")
    tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-3-8B")

    ds_config = {
        "train_batch_size": 256,
        "gradient_accumulation_steps": 4,
        "fp16": {"enabled": True},
        "zero_optimization": {
            "stage": 3,
            "offload_optimizer": {"device": "cpu"},
            "offload_param": {"device": "cpu"},
            "overlap_comm": True,
            "contiguous_gradients": True,
        },
        "optimizer": {
            "type": "AdamW",
            "params": {"lr": 2e-5, "weight_decay": 0.01}
        },
        "scheduler": {
            "type": "WarmupDecayLR",
            "params": {"warmup_num_steps": 100, "total_num_steps": 10000}
        }
    }

    model_engine, optimizer, _, _ = deepspeed.initialize(
        model=model, config=ds_config
    )

    # Training loop
    for batch in dataloader:
        loss = model_engine(batch)
        model_engine.backward(loss)
        model_engine.step()

if __name__ == "__main__":
    train()

# AWS: Deploy model with SageMaker inference
aws sagemaker create-endpoint-config \
  --endpoint-config-name "llm-endpoint-config" \
  --production-variants '[{
    "VariantName": "primary",
    "ModelName": "my-llm-model",
    "InstanceType": "ml.g5.2xlarge",
    "InitialInstanceCount": 2,
    "InitialVariantWeight": 1
  }]'

# GCP: Deploy with Vertex AI
gcloud ai endpoints deploy-model ENDPOINT_ID \
  --model=MODEL_ID \
  --display-name="llm-deployment" \
  --machine-type=g2-standard-8 \
  --accelerator=type=nvidia-l4,count=1 \
  --min-replica-count=1 \
  --max-replica-count=10

# Azure: Deploy with Azure ML
az ml online-deployment create \
  --name gpu-deployment \
  --endpoint-name llm-endpoint \
  --model azureml:my-model@latest \
  --instance-type Standard_NC24ads_A100_v4 \
  --instance-count 2