Usar o FSDP para ajustar o Llama 4 em um cluster Slurm A4

Neste tutorial, mostramos como ajustar um modelo de linguagem grande (LLM) Llama-4-Scout-17 em um cluster do Slurm com vários nós e várias GPUs no Google Cloud. O cluster usa duas instâncias de máquina virtual (VM) A4, cada uma com oito GPUs NVIDIA B200.

Os dois principais processos descritos neste tutorial são os seguintes:

  1. Implante um cluster Slurm de alto desempenho e de nível de produção usando o Google Cloud Cluster Toolkit. Como parte dessa implantação, você cria uma imagem de VM personalizada com o software necessário pré-instalado. Você também configura uma instância compartilhada do Filestore e configura redes RDMA de alta velocidade.
  2. Depois que o cluster for implantado, execute um job de ajuste refinado distribuído usando o conjunto de scripts que acompanham este tutorial. O job usa o paralelismo de dados totalmente fragmentados (FSDP) do PyTorch, que você acessa pelo aprendizado por reforço de transformadores do Hugging Face

Este tutorial é destinado a engenheiros de machine learning (ML), administradores e operadores de plataforma e especialistas em dados e IA interessados em usar os recursos de programação de jobs do Slurm para lidar com cargas de trabalho de ajuste refinado.

Objetivos

  • Acessar o Llama 4 usando o Hugging Face

  • Preparar o ambiente

  • Crie e implante um cluster Slurm A4 High-GPU de nível de produção.

  • Configure um ambiente de vários nós para treinamento distribuído com FSDP.

  • Ajuste o modelo Llama 4 usando o Hugging Face trl.SFTTrainer.

  • Organize dados em SSDs locais.

  • Monitorar o job.

  • Fazer a limpeza.

Custos

Neste documento, você vai usar os seguintes componentes faturáveis do Google Cloud:

Para gerar uma estimativa de custo baseada na projeção de uso deste tutorial, use a calculadora de preços.

Novos usuários do Google Cloud podem estar qualificados para um teste sem custo financeiro.

Antes de começar

  1. Faça login na sua conta do Google Cloud . Se você começou a usar o Google Cloud, crie uma conta para avaliar o desempenho de nossos produtos em situações reais. Clientes novos também recebem US$ 300 em créditos para executar, testar e implantar cargas de trabalho.

  2. Instale a CLI do Google Cloud.

  3. Ao usar um provedor de identidade (IdP) externo, primeiro faça login na gcloud CLI com sua identidade federada.

  4. Para inicializar a gcloud CLI, execute o seguinte comando: 

    gcloud init

  5. Crie ou selecione um Google Cloud projeto.

    Funções necessárias para selecionar ou criar um projeto

    • Selecionar um projeto: não é necessário um papel específico do IAM para selecionar um projeto. Você pode escolher qualquer projeto em que tenha recebido um papel.
    • Criar um projeto: para criar um projeto, é necessário ter o papel de Criador de projetos (roles/resourcemanager.projectCreator), que contém a permissão resourcemanager.projects.create. Saiba como conceder papéis.

    • Crie um projeto do Google Cloud :

      gcloud projects create PROJECT_ID

      Substitua PROJECT_ID por um nome para o projeto Google Cloud que você está criando.

    • Selecione o projeto Google Cloud que você criou:

      gcloud config set project PROJECT_ID

      Substitua PROJECT_ID pelo nome do projeto do Google Cloud .

  6. Verifique se o faturamento está ativado para o projeto do Google Cloud .

  7. Ative a API necessária:

    Funções necessárias para ativar APIs

    Para ativar as APIs, é necessário ter o papel do IAM de administrador do Service Usage (roles/serviceusage.serviceUsageAdmin), que contém a permissão serviceusage.services.enable. Saiba como conceder papéis. 

    gcloud services enable compute.googleapis.com file.googleapis.com logging.googleapis.com cloudresourcemanager.googleapis.com servicenetworking.googleapis.com

  8. Instale a CLI do Google Cloud.

  9. Ao usar um provedor de identidade (IdP) externo, primeiro faça login na gcloud CLI com sua identidade federada.

  10. Para inicializar a gcloud CLI, execute o seguinte comando: 

    gcloud init

  11. Crie ou selecione um Google Cloud projeto.

    Funções necessárias para selecionar ou criar um projeto

    • Selecionar um projeto: não é necessário um papel específico do IAM para selecionar um projeto. Você pode escolher qualquer projeto em que tenha recebido um papel.
    • Criar um projeto: para criar um projeto, é necessário ter o papel de Criador de projetos (roles/resourcemanager.projectCreator), que contém a permissão resourcemanager.projects.create. Saiba como conceder papéis.

    • Crie um projeto do Google Cloud :

      gcloud projects create PROJECT_ID

      Substitua PROJECT_ID por um nome para o projeto Google Cloud que você está criando.

    • Selecione o projeto Google Cloud que você criou:

      gcloud config set project PROJECT_ID

      Substitua PROJECT_ID pelo nome do projeto do Google Cloud .

  12. Verifique se o faturamento está ativado para o projeto do Google Cloud .

  13. Ative a API necessária:

    Funções necessárias para ativar APIs

    Para ativar as APIs, é necessário ter o papel do IAM de administrador do Service Usage (roles/serviceusage.serviceUsageAdmin), que contém a permissão serviceusage.services.enable. Saiba como conceder papéis. 

    gcloud services enable compute.googleapis.com file.googleapis.com logging.googleapis.com cloudresourcemanager.googleapis.com servicenetworking.googleapis.com
    14.

  14. Atribua papéis à sua conta de usuário. Execute o seguinte comando uma vez para cada um dos seguintes papéis do IAM: roles/compute.admin, roles/iam.serviceAccountUser, roles/file.editor, roles/storage.admin, roles/serviceusage.serviceUsageAdmin 

    gcloud projects add-iam-policy-binding PROJECT_ID --member="user:USER_IDENTIFIER" --role=ROLE

    Substitua:

    • PROJECT_ID: o ID do projeto.
    • USER_IDENTIFIER: o identificador da sua conta de usuário . Por exemplo, myemail@example.com.
    • ROLE: o papel do IAM concedido à sua conta de usuário.
  15. Ative a conta de serviço padrão para seu projeto do Google Cloud : 
    gcloud iam service-accounts enable PROJECT_NUMBER-compute@developer.gserviceaccount.com 
    
    --project=PROJECT_ID

    Substitua PROJECT_NUMBER pelo número do projeto. Para revisar o número do projeto, consulte Receber um projeto atual.

  16. Conceda o papel de editor (roles/editor) à conta de serviço padrão: 
    gcloud projects add-iam-policy-binding PROJECT_ID 
    
    --member="serviceAccount:PROJECT_NUMBER-compute@developer.gserviceaccount.com" 
    
    --role=roles/editor
  17. Crie as credenciais de autenticação local para sua conta de usuário: 
    gcloud auth application-default login
  18. Ative o Login do SO no seu projeto: 
    gcloud compute project-info add-metadata --metadata=enable-oslogin=TRUE
  19. Faça login ou crie uma conta do Hugging Face.
  20. Instale as dependências necessárias para usar o Cluster Toolkit.

Acessar o Llama 4 usando o Hugging Face

Para usar o Hugging Face e acessar o Llama 4, faça o seguinte:

  1. Assine o contrato de consentimento para usar o Llama 4.

  2. Crie um token de acesso read do Hugging Face.

    Clique em Seu perfil > Configurações > Tokens de acesso > +Criar novo token

  3. Copie e salve o valor do token read access. Você vai usar esse valor mais tarde neste tutorial.

Preparar o ambiente

Para preparar o ambiente, siga estas etapas:

  1. Clone o repositório do Cluster Toolkit no GitHub:

    git clone https://github.com/GoogleCloudPlatform/cluster-toolkit.git

  2. Crie um bucket do Cloud Storage:

    gcloud storage buckets create gs://BUCKET_NAME \
    --project=PROJECT_ID

    Substitua:

    • BUCKET_NAME: um nome para o bucket do Cloud Storage que segue os requisitos de nomenclatura de bucket.

    • PROJECT_ID: o ID do projetoGoogle Cloud em que você quer criar o bucket do Cloud Storage.

Criar um cluster Slurm A4

Para criar um cluster do Slurm A4, siga estas etapas:

  1. Acesse o diretório cluster-toolkit:

    cd cluster-toolkit

  2. Se esta for a primeira vez que você usa o Cluster Toolkit, crie o binário gcluster:

    make

  3. Acesse o diretório examples/machine-learning/a4-highgpu-8g:

    cd examples/machine-learning/a4-highgpu-8g/

  4. Abra o arquivo a4high-slurm-deployment.yaml e edite-o da seguinte forma:

    terraform_backend_defaults:
  type: gcs
  configuration:
    bucket: BUCKET_NAME

vars:
  deployment_name: a4-high
  project_id: PROJECT_ID
  region: REGION
  zone: ZONE
  a4h_cluster_size: 2
  a4h_reservation_name: RESERVATION_URL

    Substitua:

    • BUCKET_NAME: o nome do bucket do Cloud Storage criado na seção anterior.

    • PROJECT_ID: o ID do projetoGoogle Cloud em que o Cloud Storage existe e em que você quer criar o cluster do Slurm.

    • REGION: a região em que sua reserva existe.

    • ZONE: a zona em que sua reserva existe.

    • RESERVATION_URL: o URL da reserva que você quer usar para criar o cluster do Slurm. Com base no projeto em que a reserva existe, especifique um dos seguintes valores:

      • A reserva existe no seu projeto: RESERVATION_NAME

      • A reserva existe em um projeto diferente, e seu projeto pode usar a reserva: projects/RESERVATION_PROJECT_ID/reservations/RESERVATION_NAME

  5. Implante o cluster:

    ./gcluster deploy -d examples/machine-learning/a4-highgpu-8g/a4high-slurm-deployment.yaml examples/machine-learning/a4-highgpu-8g/a4high-slurm-blueprint.yaml --auto-approve

    O comando ./gcluster deploy é um processo de duas fases:

    • A primeira fase cria uma imagem personalizada com todo o software pré-instalado, o que pode levar até 35 minutos.

    • A segunda fase implanta o cluster usando essa imagem personalizada. Esse processo deve ser concluído mais rapidamente do que a primeira fase.

    Se a primeira fase for concluída, mas a segunda falhar, tente implantar o cluster do Slurm novamente pulando a primeira fase:

    ./gcluster deploy -d examples/machine-learning/a4-highgpu-8g/a4high-slurm-deployment.yaml examples/machine-learning/a4-highgpu-8g/a4high-slurm-blueprint.yaml --auto-approve --skip "image" -w

Preparar sua carga de trabalho

Para preparar sua carga de trabalho, faça o seguinte:

  1. Crie scripts de carga de trabalho.

  2. Faça upload dos scripts para o cluster do Slurm.

  3. Conecte-se ao cluster do Slurm.

  4. Instale frameworks e ferramentas.

Criar scripts de carga de trabalho

Para criar os scripts que sua carga de trabalho de ajuste refinado vai usar, siga estas etapas:

  1. Para configurar o ambiente virtual Python, crie o arquivo install_environment.sh com o seguinte conteúdo:

    #!/bin/bash
# This script sets up a consistent environment for FSDP training.
# It is meant to be run once on the login node of your Slurm cluster
set -e

# --- 1. Create the Python virtual environment ---
VENV_PATH="$HOME/.venv/venv-fsdp"
if [ ! -d "$VENV_PATH" ]; then
  echo "--- Creating Python virtual environment at $VENV_PATH ---"
  python3 -m venv $VENV_PATH
else
  echo "--- Virtual environment already exists at $VENV_PATH ---"
fi

source $VENV_PATH/bin/activate

# --- 2. Install Dependencies ---
echo "--- [STEP 2.1] Upgrading build toolchain ---"
pip install --upgrade pip wheel packaging

echo "--- [STEP 2.2] Installing PyTorch Nightly ---"
pip install --force-reinstall --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/cu128

echo "--- [STEP 2.3] Installing application dependencies ---"
if [ -f "requirements-fsdp.txt" ]; then
    pip install -r requirements-fsdp.txt
else
    echo "ERROR: requirements-fsdp.txt not found!"
    exit 1
fi

# --- 3. Download the Model ---
echo "--- [STEP 2.4] Downloading Llama4 model ---"
if [ -z "$HF_TOKEN" ]; then
  echo "ERROR: The HF_TOKEN environment variable is not set."; exit 1;
fi
pip install huggingface_hub[cli]

# Execute the CLI using its full, explicit path
$VENV_PATH/bin/huggingface-cli download meta-llama/Llama-4-Scout-17B-16E-Instruct --local-dir ~/Llama-4-Scout-17B-16E-Instruct --token $HF_TOKEN

echo "--- Environment setup complete. ---"

    Esse script configura um ambiente virtual Python confiável, instala uma versão noturna do PyTorch e baixa o modelo Llama 4.

  2. Para especificar as dependências do Python para o script de treinamento, crie um arquivo requirements-fsdp.txt com o seguinte conteúdo:

    transformers==4.55.0
datasets==4.0.0
peft==0.16.0
accelerate==1.9.0
trl==0.21.0

# Other dependencies
sentencepiece==0.2.0

  3. Especifique llama4-train-distributed.py como o script de treinamento principal:

    import torch
from datasets import load_dataset
from peft import LoraConfig, PeftModel
from transformers import (
    AutoModelForCausalLM,
    AutoTokenizer,
    TrainingArguments,
    HfArgumentParser,
)

from torch.distributed import get_rank, get_world_size

from transformers.models.llama4.modeling_llama4 import Llama4TextDecoderLayer
from trl import SFTTrainer
from dataclasses import dataclass, field
from typing import Optional

@dataclass
class ScriptArguments:
    model_id: str = field(metadata={"help": "Hugging Face model ID from the Hub"})
    dataset_name: str = field(default="philschmid/gretel-synthetic-text-to-sql", metadata={"help": "Dataset from the Hub"})
    run_inference_after_training: bool = field(default=False, metadata={"help": "Run sample inference on rank 0 after training"})
    dataset_subset_size: Optional[int] = field(default=None, metadata={"help": "Number of samples to use from the dataset for training. If None, uses the full dataset."})

@dataclass
class PeftArguments:
    lora_r: int = field(default=16, metadata={"help": "LoRA attention dimension"})
    lora_alpha: int = field(default=32, metadata={"help": "LoRA alpha scaling factor"})
    lora_dropout: float = field(default=0.05, metadata={"help": "LoRA dropout probability"})

@dataclass
class SftTrainingArguments(TrainingArguments):
    max_length: Optional[int] = field(default=2048, metadata={"help": "The maximum sequence length for SFTTrainer"})
    packing: Optional[bool] = field(default=False, metadata={"help": "Enable packing for SFTTrainer"})
    ddp_find_unused_parameters: Optional[bool] = field(default=True, metadata={"help": "When using FSDP activation checkpointing, this must be set to True"})

def formatting_prompts_func(example):
    system_message = "You are a text to SQL query translator. Users will ask you questions in English and you will generate a SQL query based on the provided SCHEMA."
    user_prompt = f"### SCHEMA:\n{example['sql_context']}\n\n### USER QUERY:\n{example['sql_prompt']}"
    response = f"\n\n### SQL QUERY:\n{example['sql']}"
    return f"{system_message}\n\n{user_prompt}{response}"

def main():
    parser = HfArgumentParser((ScriptArguments, PeftArguments, SftTrainingArguments))
    script_args, peft_args, training_args = parser.parse_args_into_dataclasses()

    training_args.gradient_checkpointing = True
    training_args.gradient_checkpointing_kwargs = {"use_reentrant": False}

    training_args.optim = "adamw_torch_fused"

    training_args.fsdp = "full_shard"
    training_args.fsdp_config = {
        "fsdp_auto_wrap_policy": "TRANSFORMER_BASED_WRAP",
        "fsdp_transformer_layer_cls_to_wrap": [Llama4TextDecoderLayer],
        "fsdp_state_dict_type": "FULL_STATE_DICT",
        "fsdp_offload_params": False,
        "fsdp_forward_prefetch": True,
    }

    tokenizer = AutoTokenizer.from_pretrained(script_args.model_id, trust_remote_code=True)

    model = AutoModelForCausalLM.from_pretrained(
        script_args.model_id,
        torch_dtype=torch.bfloat16,
        trust_remote_code=True,
        attn_implementation="sdpa",
    )

    peft_config = LoraConfig(
        r=peft_args.lora_r,
        lora_alpha=peft_args.lora_alpha,
        lora_dropout=peft_args.lora_dropout,
        bias="none",
        task_type="CAUSAL_LM",
        target_modules=["q_proj", "v_proj", "k_proj", "o_proj", "gate_proj", "up_proj", "down_proj"],
    )
    rank = get_rank()
    world_size = get_world_size()

    dataset = load_dataset(script_args.dataset_name, split="train")

    if script_args.dataset_subset_size is not None:
        dataset = dataset.select(range(script_args.dataset_subset_size))
    else:
        print(f"Using the full dataset with {len(dataset)} samples.")

    dataset = dataset.shuffle(seed=training_args.seed)
    print(f"Dataset shuffled with seed: {training_args.seed}.")

    if world_size > 1:
        print(f"Sharding dataset for Rank {rank} of {world_size}.")
        dataset = dataset.shard(num_shards=world_size, index=rank)

    print("Initializing SFTTrainer...")
    trainer = SFTTrainer(
        model=model,
        args=training_args,
        train_dataset=dataset,
        peft_config=peft_config,
        formatting_func=formatting_prompts_func,
        processing_class=tokenizer,
    )

    trainer.train()

    trainer.save_model(training_args.output_dir)

    if script_args.run_inference_after_training and trainer.is_world_process_zero():
        del model
        del trainer
        torch.cuda.empty_cache()
        run_post_training_inference(script_args, training_args, tokenizer)

def run_post_training_inference(script_args, training_args, tokenizer):
    """
    Loads the fine-tuned PEFT adapter from the local output directory and runs inference.
    This should only be called on rank 0 after training is complete.
    """
    print("\n" + "="*50)
    print("=== RUNNING POST-TRAINING INFERENCE TEST ===")
    print("="*50 + "\n")

    # Load the base model and merge the adapter.
    base_model = AutoModelForCausalLM.from_pretrained(
        script_args.model_id,
        torch_dtype=torch.bfloat16,
        trust_remote_code=True,
        device_map="auto"
    )
    # Load the PEFT adapter and merge it into the base model
    model = PeftModel.from_pretrained(base_model, training_args.output_dir)
    model = model.merge_and_unload() # Merge weights for faster inference
    model.eval()

    # Define the test case
    schema = "CREATE TABLE artists (Name TEXT, Country TEXT, Genre TEXT)"
    system_message = "You are a text to SQL query translator. Users will ask you questions in English and you will generate a SQL query based on the provided SCHEMA."
    question = "Show me all artists from the Country just north of the USA."

    # This must match the formatting_func exactly
    prompt = f"{system_message}\n\n### SCHEMA:\n{schema}\n\n### USER QUERY:\n{question}\n\n### SQL QUERY:\n"

    print(f"Test Prompt:\n{prompt}")

    inputs = tokenizer(prompt, return_tensors="pt").to("cuda")

    print("\n--- Generating SQL... ---")
    outputs = model.generate(
        **inputs,
        max_new_tokens=100,
        pad_token_id=tokenizer.eos_token_id,
        do_sample=False,
        temperature=None,
        top_p=None,
    )

    generated_sql = tokenizer.decode(outputs[0], skip_special_tokens=True)[len(prompt):].strip()

    print(f"\n--- Generated SQL Query ---")
    print(generated_sql)
    print("\n" + "="*50)
    print("=== INFERENCE TEST COMPLETE ===")
    print("="*50 + "\n")

if __name__ == "__main__":
    main()

    Esse script usa o treinador de ajuste supervisionado (SFT) da TRL para gerenciar loops de treinamento FSDP, configuração de adaptação de classificação baixa (LoRA) e formatação de dados.

  4. Para especificar as tarefas que os jobs vão executar no cluster Slurm, crie o arquivo submit.slurm com o seguinte conteúdo:

    #!/bin/bash
#SBATCH --job-name=llama4-fsdp-fixed
#SBATCH --nodes=2
#SBATCH --ntasks-per-node=8
#SBATCH --gpus-per-node=8
#SBATCH --partition=a4high
#SBATCH --output=llama4-%j.out
#SBATCH --error=llama4-%j.err

set -e
set -x

echo "--- Slurm Job Started ---"
echo "Job ID: $SLURM_JOB_ID"
echo "Node List: $SLURM_JOB_NODELIST"

# --- Define Paths ---
LOCAL_SSD_PATH="/mnt/localssd/job_${SLURM_JOB_ID}"
VENV_PATH="${HOME}/.venv/venv-fsdp"
MODEL_PATH="${HOME}/Llama-4-Scout-17B-16E-Instruct"

# --- STAGE 1: Stage Data to Local SSD on Each Node ---
srun --ntasks=$SLURM_NNODES --ntasks-per-node=1 bash -c "
  echo '--- Staging on node: $(hostname) ---'

  mkdir -p ${LOCAL_SSD_PATH}

  echo 'Copying virtual environment...'
  rsync -a -q ${VENV_PATH}/ ${LOCAL_SSD_PATH}/venv/

  echo 'Copying model weights...'
  rsync -a --info=progress2 ${MODEL_PATH}/ ${LOCAL_SSD_PATH}/model/

  mkdir -p ${LOCAL_SSD_PATH}/hf_cache

  echo '--- Staging on $(hostname) complete ---'
"
echo "--- Staging complete on all nodes ---"

# --- STAGE 2: Run the Training Job ---
echo "--- Launching Distributed Training with GIB NCCL Plugin ---"
nodes=( $( scontrol show hostnames "$SLURM_JOB_NODELIST" ) )
head_node=${nodes[0]}
head_node_ip=$(srun --nodes=1 --ntasks=1 -w "$head_node" hostname --ip-address)

export MASTER_ADDR=$head_node_ip
export MASTER_PORT=29500

export NCCL_SOCKET_IFNAME=enp0s19
export NCCL_NET=gIB

# export NCCL_DEBUG=INFO # Un-comment to diagnose NCCL issues if needed

srun --cpu-bind=none --accel-bind=g bash -c '
  # Activate the environment from the local copy
  source '${LOCAL_SSD_PATH}'/venv/bin/activate

  # Point Hugging Face cache to the local SSD
  export HF_HOME='${LOCAL_SSD_PATH}'/hf_cache

  export RANK=$SLURM_PROCID
  export WORLD_SIZE=$SLURM_NTASKS
  export LOCAL_RANK=$SLURM_LOCALID

  export LD_LIBRARY_PATH=/usr/local/gib/lib64:$LD_LIBRARY_PATH
  source /usr/local/gib/scripts/set_nccl_env.sh

  # --- Launch the training ---
  python \
    '${SLURM_SUBMIT_DIR}'/llama4-train-distributed.py \
      --model_id="'${LOCAL_SSD_PATH}'/model/" \
      --output_dir="'${LOCAL_SSD_PATH}'/outputs/" \
      --dataset_name="philschmid/gretel-synthetic-text-to-sql" \
      --seed=900913 \
      --bf16=True \
      --num_train_epochs=1 \
      --per_device_train_batch_size=2 \
      --gradient_accumulation_steps=4 \
      --learning_rate=2e-5 \
      --logging_steps=10 \
      --lora_r=16 \
      --lora_alpha=32 \
      --lora_dropout=0.05 \
      --run_inference_after_training
'

# --- STAGE 3: Copy Final Results Back to Persistent Storage ---
echo "--- Copying final results from local SSD to shared storage ---"
PERSISTENT_OUTPUT_DIR="${HOME}/outputs/llama4_job_${SLURM_JOB_ID}"
mkdir -p "$PERSISTENT_OUTPUT_DIR"

# Only copy from the head node where trl has combined the results
srun --nodes=1 --ntasks=1 -w "$head_node" \
  rsync -a --info=progress2 "${LOCAL_SSD_PATH}/outputs/" "${PERSISTENT_OUTPUT_DIR}/"

# --- STAGE 4: Cleanup ---
echo "--- Cleaning up local SSD on all nodes ---"
srun --ntasks=$SLURM_NNODES --ntasks-per-node=1 bash -c "rm -rf ${LOCAL_SSD_PATH}"

echo "--- Slurm Job Finished ---"

Fazer upload de scripts para o cluster Slurm

Para fazer upload dos scripts criados na seção anterior para o cluster do Slurm, siga estas etapas:

  1. Para identificar o nó de login, liste todas as VMs A4 no projeto:

    gcloud compute instances list --filter="machineType:a4-highgpu-8g"

    O nome do nó de login é semelhante a a4-high-login-001.

  2. Faça upload dos scripts para o diretório principal do nó de login:

    gcloud compute scp --project="$PROJECT_ID" --zone="$ZONE" --tunnel-through-iap \
  ./install_environment.sh \
  ./requirements-fsdp.txt \
  ./llama4-train-distributed.py \
  ./submit.slurm \
  "${LOGIN_NODE_NAME}":~/

    Substitua LOGIN_NODE_NAME pelo nome do nó de login.

Conectar-se ao cluster do Slurm

Conecte-se ao cluster do Slurm conectando-se ao nó de login por SSH:

gcloud compute ssh LOGIN_NODE_NAME \
    --project=PROJECT_ID \
    --tunnel-through-iap \
    --zone=ZONE

Instalar frameworks e ferramentas

Depois de se conectar ao nó de login, instale frameworks e ferramentas fazendo o seguinte:

  1. Exporte seu token do Hugging Face:

    # On the login node
export HF_TOKEN="hf_..." # Replace with your token

  2. Execute o script de instalação:

    # On the login node
chmod +x install_environment.sh
./install_environment.sh

    Esse comando configura um ambiente virtual com todas as dependências necessárias e faz o download dos pesos do modelo no arquivo ~/Llama-4-Scout-17B-16E-Instruct.

    Como o download do modelo é muito grande (~200 GB), esse processo leva cerca de 30 minutos, dependendo das condições da rede.

Iniciar sua carga de trabalho de ajuste refinado

Para começar a treinar sua carga de trabalho, faça o seguinte:

  1. Envie o job para o programador do Slurm:

    sbatch submit.slurm

  2. No nó de login do cluster do Slurm, monitore o progresso do job verificando os arquivos de saída criados no diretório home:

    # On the login node
tail -f llama4-*.out

    Se o job for iniciado com sucesso, o arquivo .err vai mostrar uma barra de progresso que é atualizada conforme o job avança.

    Esse job leva pouco mais de uma hora para ser concluído no cluster do Slurm. O job tem duas fases principais:

    • Copiar o modelo de base grande para o SSD local de cada nó de computação.
    • O job de treinamento, que começa quando a cópia do modelo é concluída. Esse job leva cerca de 35 minutos para ser executado.

Limpar

Para evitar cobranças na sua conta do Google Cloud pelos recursos usados no tutorial, exclua o projeto que os contém ou mantenha o projeto e exclua os recursos individuais.

Excluir o projeto

Excluir um projeto do Google Cloud :

gcloud projects delete PROJECT_ID

Excluir o cluster do Slurm

Para excluir o cluster do Slurm, siga estas etapas:

  1. Acesse o diretório cluster-toolkit.

  2. Destrua o arquivo do Terraform e todos os recursos criados:

    ./gcluster destroy a4-high --auto-approve

Excluir a instância do Filestore

Por padrão, a instância do Filestore tem a configuração deletion_protection definida como "true" no blueprint cluster-toolkit. Essa configuração evita a perda acidental de dados ao modificar ambientes. Para excluir a instância do Filestore, desative manualmente a proteção contra exclusão.

A seguir