O Google usa tecnologia de IA na tradução de conteúdos para seu idioma de preferência. As traduções com IA podem ter erros.

Usar o FSDP para ajustar o Mixtral-8x7B em um cluster A4 Slurm

GPUs em cluster

Neste tutorial, você vai aprender a ajustar um modelo mistralai/Mixtral-8x7B-v0.1 em um cluster Slurm de vários nós e várias GPUs no Google Cloud. O cluster usa duas instâncias de máquina virtual (VM) a4-highgpu-8g, cada uma com 8 GPUs NVIDIA B200.

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

Implante um cluster Slurm de alta performance e 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 um sistema de arquivos Lustre compartilhado e uma rede de alta velocidade.
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, na sigla em inglês) do PyTorch, que você acessa pela biblioteca Transformer Reinforcement Learning (TRL) do Hugging Face.

Este tutorial é destinado a engenheiros de machine learning (ML), pesquisadores, administradores e operadores de plataforma e especialistas em dados e IA que têm interesse em distribuir uma carga de trabalho de IA em vários nós e GPUs.

Objetivos

Acessar o Mixtral usando o Hugging Face
Instalar o Cluster Toolkit
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 Mixtral usando a classe trl.SFTTrainer do Hugging Face.
Faça o stage dos dados para 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

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.

Instale a CLI do Google Cloud.

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

Para inicializar a CLI gcloud, execute o seguinte comando:

gcloud init

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 .

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

Ative a API necessária:

Funções necessárias para ativar APIs

Para ativar APIs, você precisa da permissão serviceusage.services.enable. Se você criou o projeto, provavelmente já tem essa permissão com o papel de Proprietário (roles/owner). Caso contrário, é possível receber essa permissão com o papel de Administrador do Service Usage (roles/serviceusage.serviceUsageAdmin). Saiba como conceder papéis.

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

Instale a CLI do Google Cloud.

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

Para inicializar a CLI gcloud, execute o seguinte comando:

gcloud init

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 .

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

Ative a API necessária:

Funções necessárias para ativar APIs

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

Atribua papéis à sua conta de usuário. Execute uma vez o seguinte comando 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.
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.

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

Crie as credenciais de autenticação local para sua conta de usuário:
```
gcloud auth application-default login
```
Observação: se você usa um shell local e um provedor de identidade (IdP) externo e encontrar um erro de autenticação depois de executar o comando anterior, faça login na Google Cloud CLI com sua identidade federada.

Ative o Login do SO no seu projeto:

gcloud compute project-info add-metadata --metadata=enable-oslogin=TRUE

Faça login ou crie uma conta do Hugging Face.
Instale as dependências necessárias para usar o Cluster Toolkit.

Acessar o Mixtral usando o Hugging Face

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

Faça login no Hugging Face e conheça o modelo Mixtral.
Crie um token de acesso read do Hugging Face.
Copie e salve o valor do token. Você vai usá-lo mais tarde neste tutorial.

Instalar o Cluster Toolkit

O Cluster Toolkit é uma ferramenta de código aberto que simplifica a implantação de cargas de trabalho de computação de alto desempenho (HPC), inteligência artificial (IA) e machine learning (ML) no Google Cloud. Para mais informações sobre como usar gcluster e gerenciar clusters, consulte a visão geral do Cluster Toolkit.

Prepare a versão do Cluster Toolkit:

export CLUSTER_TOOLKIT_TAG=v1.97.0

# Detect OS (linux or mac)
case "$(uname -s)" in
  Linux*)     OS="linux" ;;
  Darwin*)    OS="mac" ;;
  *)          echo "Error: Unsupported operating system: $(uname -s)" >&2; exit 1 ;;
esac

# Detect Architecture (amd64 or arm64)
case "$(uname -m)" in
  x86_64)     ARCH="amd64" ;;
  aarch64|arm64) ARCH="arm64" ;;
  *)          echo "Error: Unsupported architecture: $(uname -m)" >&2; exit 1 ;;
esac

Faça o download da versão:

# Download and extract the platform-specific bundle
curl -LO "https://github.com/GoogleCloudPlatform/cluster-toolkit/releases/download/${CLUSTER_TOOLKIT_TAG}/gcluster_bundle_${OS}_${ARCH}.zip"
unzip "gcluster_bundle_${OS}_${ARCH}.zip" -d cluster-toolkit/
rm -f "gcluster_bundle_${OS}_${ARCH}.zip"

Defina o caminho gcluster:

export CLUSTER_TOOLKIT_PATH="$(pwd)/cluster-toolkit"
export PATH="${CLUSTER_TOOLKIT_PATH}:${PATH}"
gcluster --version

Preparar o ambiente

Para preparar o ambiente, siga estas etapas:

Defina as variáveis de ambiente padrão:
```
export PROJECT_ID="YOUR_PROJECT_ID"
export ZONE="YOUR_ZONE"
export REGION="YOUR_REGION"
export RESERVATION_NAME="YOUR_RESERVATION_NAME"
export DEPLOYMENT_NAME="YOUR_CLUSTER_NAME"
export BUCKET_NAME="YOUR_GCS_BUCKET"
export HF_TOKEN="YOUR_HF_TOKEN"
```
Substitua:
- YOUR_PROJECT_ID: o nome do projeto Google Cloud em que você quer criar o cluster do Slurm.
- YOUR_ZONE: a zona em que a reserva existe.
- YOUR_REGION: a região em que sua reserva existe.
- YOUR_RESERVATION_NAME: o URL ou o nome da reserva que você quer usar para criar o cluster do Slurm.
- YOUR_CLUSTER_NAME: o nome da sua implantação. Use um nome curto que contenha apenas letras e números (por exemplo, a4high). Esse nome também é atribuído ao cluster do Slurm criado pela implantação.
- YOUR_GCS_BUCKET: o nome do bucket em que você armazena os resultados do checkpoint de treinamento. Especifique um bucket atual ou crie um novo. Antes de criar, familiarize-se com os requisitos de nomenclatura de bucket.
- YOUR_HF_TOKEN: o token do Hugging Face que você criou em uma etapa anterior.

Crie um bucket do Cloud Storage:

gcloud storage buckets create "gs://${BUCKET_NAME}" \
  --project="${PROJECT_ID}"

Criar um cluster Slurm A4

Para criar um cluster do Slurm A4, faça o seguinte:

Substitua o arquivo a4high-slurm-deployment.yaml no diretório examples/machine-learning/a4-highgpu-8g pelas variáveis de ambiente executando o seguinte comando:

MANIFEST_PATH="${CLUSTER_TOOLKIT_PATH}/examples/machine-learning/a4-highgpu-8g"
cat <<EOF > "${MANIFEST_PATH}/a4high-slurm-deployment.yaml"
terraform_backend_defaults:
  type: gcs
  configuration:
    bucket: ${BUCKET_NAME}

vars:
  deployment_name: ${DEPLOYMENT_NAME}
  project_id: ${PROJECT_ID}
  region: ${REGION}
  zone: ${ZONE}
  a4h_cluster_size: 2
  a4h_reservation_name: ${RESERVATION_NAME}
EOF

Abra o arquivo a4high-slurm-blueprint.yaml no diretório examples/machine-learning/a4-highgpu-8g e edite-o para usar o Managed Lustre no diretório compartilhado /home da seguinte maneira:
- Remova o bloco de módulo homefs padrão com source: modules/file-system/filestore.
- Ative os módulos lustrefs (bloco homefs com remote_mount: lustrefs) e private-service-access.
- No bloco vars, configure o seguinte:
  1. Mude o valor de install_managed_lustre em /var/tmp/slurm_vars.json para true.
  2. Defina o parâmetro per_unit_storage_throughput como 500.
  3. Defina o parâmetro lustre_size_gib como 36000.
  4. Remova o comentário de lustre_instance_id: lustre-instance.
  5. Comente ou remova filestore_ip_range não utilizado.
Implante o cluster:
```
./gcluster deploy \
  -d "${MANIFEST_PATH}/a4high-slurm-deployment.yaml" \
  "${MANIFEST_PATH}/a4high-slurm-blueprint.yaml" \
  --auto-approve
```
O comando ./gcluster deploy inicia um processo de duas fases, que é o seguinte:
- 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.

Preparar sua carga de trabalho

Para preparar sua carga de trabalho, siga estas etapas:

Crie scripts de carga de trabalho.
Faça upload dos scripts para o cluster do Slurm.
Conecte-se ao cluster do Slurm .
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:

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

#!/bin/bash
# This script sets a reliable environment for FSDP training.
# It is meant to be run on a compute node.
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 --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

# --- [STEP 2.4] Build Flash Attention from Source ---
echo "--- Building flash-attn from source... This will take a while. ---"
# Bypass PyTorch CUDA version mismatch check for extension builds
python3 -c "
import torch.utils.cpp_extension as ce
p = ce.__file__
content = open(p).read().replace('raise RuntimeError(CUDA_MISMATCH_MESSAGE', 'pass # raise RuntimeError')
open(p, 'w').write(content)
"

# Use all available CPU cores to speed up the build
MAX_JOBS=$(nproc) pip install flash-attn --no-build-isolation

# --- 3. Download the Model ---
echo "--- [STEP 2.5] Downloading Mixtral 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/hf download mistralai/Mixtral-8x7B-v0.1 --local-dir ~/Mixtral-8x7B-v0.1 --token $HF_TOKEN

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

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
protobuf==6.31.1

Especifique train-mixtral.py como o script de treinamento principal:

import torch
from torch.distributed.fsdp import MixedPrecision
from datasets import load_dataset
import shutil
import os
import torch.distributed as dist

from peft import LoraConfig, PeftModel, get_peft_model
from transformers import (
    AutoModelForCausalLM,
    AutoTokenizer,
    TrainingArguments,
    HfArgumentParser,
)

from torch.distributed import get_rank, get_world_size

from transformers.models.mixtral.modeling_mixtral import MixtralDecoderLayer
from trl import SFTTrainer
from dataclasses import dataclass, field
from typing import Optional

@dataclass
class ScriptArguments:
    model_id: str = field(default="mistralai/Mixtral-8x7B-v0.1", 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=False, metadata={"help": "When using FSDP activation checkpointing, this must be set to False for Mixtral"})

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": True}

    training_args.optim = "adamw_torch_fused"

    bf16_policy = MixedPrecision(
        param_dtype=torch.bfloat16,
        reduce_dtype=torch.bfloat16,
        buffer_dtype=torch.bfloat16,
    )

    training_args.fsdp = "full_shard"
    training_args.fsdp_config = {
        "fsdp_auto_wrap_policy": "TRANSFORMER_BASED_WRAP",
        "fsdp_transformer_layer_cls_to_wrap": [MixtralDecoderLayer],
        "fsdp_state_dict_type": "SHARDED_STATE_DICT",
        "fsdp_offload_params": False,
        "fsdp_forward_prefetch": True,
        "fsdp_mixed_precision_policy": bf16_policy
    }

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

    tokenizer.pad_token = tokenizer.eos_token
    tokenizer.padding_side = "right"

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

    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"],
    )

    model = get_peft_model(model, peft_config)

    data_splits = load_dataset(script_args.dataset_name)

    dataset = data_splits["train"]
    eval_dataset = data_splits["test"]

    if script_args.dataset_subset_size is not None:
        dataset = dataset.select(range(script_args.dataset_subset_size))

    dataset = dataset.shuffle(seed=training_args.seed)

    trainer = SFTTrainer(
        model=model,
        args=training_args,
        train_dataset=dataset,
        eval_dataset=eval_dataset,
        formatting_func=formatting_prompts_func,
        processing_class=tokenizer,
    )

    trainer.train()

    dist.barrier()
    if trainer.is_world_process_zero():
        best_model_path = trainer.state.best_model_checkpoint

        final_model_dir = os.path.join(training_args.output_dir, "final_best_model")
        print(f"Copying best model to: {final_model_dir}")

        if os.path.exists(final_model_dir):
            shutil.rmtree(final_model_dir)
        shutil.copytree(best_model_path, final_model_dir)

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

def run_post_training_inference(script_args, best_model_path, tokenizer):
    print("\n" + "="*50)
    print("=== RUNNING POST-TRAINING INFERENCE TEST ===")
    print("="*50 + "\n")

    base_model = AutoModelForCausalLM.from_pretrained(
        script_args.model_id,
        torch_dtype=torch.bfloat16,
        trust_remote_code=True,
        device_map="auto"
    )
    model = PeftModel.from_pretrained(base_model, best_model_path)
    model = model.merge_and_unload()
    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."

    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()

Para especificar as tarefas que os jobs vão executar no cluster Slurm, crie o arquivo train-mixtral.sh com o seguinte conteúdo:

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

set -e
set -x

echo "--- Slurm Job Started ---"

# --- Define Paths ---
LOCAL_SSD_PATH="/mnt/localssd/job_${SLURM_JOB_ID}"
VENV_PATH="${HOME}/.venv/venv-fsdp"
MODEL_PATH="${HOME}/Mixtral-8x7B-v0.1"

# --- 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 ${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}'/train-mixtral.py \
    --model_id="'${LOCAL_SSD_PATH}'/model/" \
    --output_dir="${HOME}/outputs/mixtral_job_${SLURM_JOB_ID}" \
    --dataset_name="philschmid/gretel-synthetic-text-to-sql" \
    --seed=900913 \
    --bf16=True \
    --num_train_epochs=3 \
    --per_device_train_batch_size=32 \
    --gradient_accumulation_steps=4 \
    --learning_rate=4e-5 \
    --logging_steps=3 \
    --lora_r=32 \
    --lora_alpha=32 \
    --lora_dropout=0.05 \
    --eval_strategy=steps \
    --eval_steps=10 \
    --save_strategy=steps \
    --save_steps=10 \
    --load_best_model_at_end=False \
    --metric_for_best_model=eval_loss \
    --run_inference_after_training \
    --dataset_subset_size=67000
'

# --- STAGE 3: 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 dos scripts para o cluster do Slurm

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

Defina a variável LOGIN_NODE recuperando o nome do nó de login do cluster:

LOGIN_NODE="$(gcloud compute instances list \
                --project="${PROJECT_ID}" \
                --filter="labels.ghpc_deployment='${DEPLOYMENT_NAME}' AND labels.slurm_instance_role='login'" \
                --format="value(name)" | head -n 1)"

A variável LOGIN_NODE armazena um valor semelhante a ${DEPLOYMENT_NAME}-login-001.

Crie uma regra de firewall:

gcloud compute firewall-rules create allow-ssh-ingress-from-iap \
  --project="${PROJECT_ID}" \
  --network="${CLUSTER_NETWORK}" \
  --direction=INGRESS \
  --action=allow \
  --rules=tcp:22 \
  --source-ranges=35.235.240.0/20 \
  --description="Allow SSH ingress from Google Cloud Identity-Aware Proxy (IAP)"

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 \
  ./train-mixtral.py \
  ./train-mixtral.sh \
  "${LOGIN_NODE}":~/

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}" \
    --project="${PROJECT_ID}" \
    --tunnel-through-iap \
    --zone="${ZONE}"
    -- -t "export HF_TOKEN='${HF_TOKEN}'; bash -l"

Instalar frameworks e ferramentas

Depois de se conectar ao nó de login, instale as ferramentas e os frameworks necessários executando o script de instalação em um nó de computação:

# On the login node
srun \
  --job-name=env-setup \
  --nodes=1 \
  --ntasks=1 \
  --gpus-per-node=1 \
  --partition=a4high \
  bash ./install_environment.sh

Esse comando configura o ambiente virtual, instala todas as dependências e baixa as ponderações do modelo Mixtral em ~/Mixtral-8x7B-v0.1. Esse processo pode levar mais de 30 minutos.

Executar esta etapa via srun delega o script de instalação do nó de login para um nó de computação alocado, permitindo que o script acesse os drivers de GPU durante a compilação.

Iniciar a carga de trabalho de ajuste refinado

Para iniciar a carga de trabalho de treinamento, faça o seguinte:

Envie o job para o programador do Slurm:

# On the login node
sbatch train-mixtral.sh

Monitore sua carga de trabalho

Para monitorar o progresso do seu job de ajuste refinado, use os seguintes métodos:

No nó de login do cluster do Slurm, é possível monitorar o progresso do job verificando os arquivos de saída criados no diretório home:
```
# On the login node:
tail -f mixtral-*.out mixtral-*.err
```
Se o job for iniciado com sucesso, o arquivo .err vai mostrar uma barra de progresso que é atualizada conforme o job avança.

O trabalho 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.
O job inteiro leva cerca de 60 minutos para ser executado.

Também é possível monitorar o uso das GPUs no cluster do Slurm para verificar se o job de ajuste refinado está sendo executado de maneira eficiente. Para fazer isso, abra o seguinte link no navegador:

https://console.cloud.google.com/monitoring/metrics-explorer?project=YOUR_PROJECT_ID&pageState=%7B%22xyChart%22%3A%7B%22dataSets%22%3A%5B%7B%22timeSeriesFilter%22%3A%7B%22filter%22%3A%22metric.type%3D%5C%22agent.googleapis.com%2Fgpu%2Futilization%5C%22%20resource.type%3D%5C%22gce_instance%5C%22%22%2C%22perSeriesAligner%22%3A%22ALIGN_MEAN%22%7D%2C%22plotType%22%3A%22LINE%22%7D%5D%7D%7D

Como alternativa, insira o comando diretamente no terminal:

open "https://console.cloud.google.com/monitoring/metrics-explorer?project=${PROJECT_ID}&pageState=%7B%22xyChart%22%3A%7B%22dataSets%22%3A%5B%7B%22timeSeriesFilter%22%3A%7B%22filter%22%3A%22metric.type%3D%5C%22agent.googleapis.com%2Fgpu%2Futilization%5C%22%20resource.type%3D%5C%22gce_instance%5C%22%22%2C%22perSeriesAligner%22%3A%22ALIGN_MEAN%22%7D%2C%22plotType%22%3A%22LINE%22%7D%5D%7D%7D"

Ao monitorar sua carga de trabalho, você pode ver o seguinte:

Uso da GPU: para um job de ajuste refinado saudável, é esperado que o uso de todas as 16 GPUs (oito GPUs para cada VM no cluster) aumente e se estabilize em um nível específico durante todo o treinamento.
Duração do job: o job leva aproximadamente uma hora para ser concluído.

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 os recursos

Execute os comandos a seguir para excluir os recursos criados neste tutorial:

Para excluir o cluster Slurm, acesse o diretório cluster-toolkit e execute o seguinte comando:
```
./gcluster destroy "${DEPLOYMENT_NAME}" --auto-approve --robust
```

Exclua o bucket:

gcloud storage rm --recursive "gs://${BUCKET_NAME}" --quiet || true

Para excluir uma imagem do Packer, abra o navegador da Web, acesse a página a seguir, pesquise sua imagem específica e clique em "Excluir".
```
http://console.cloud.google.com/compute/images
```

Excluir o projeto

Excluir um projeto do Google Cloud :

gcloud projects delete PROJECT_ID

A seguir

Reimplantar um cluster do Slurm
Testar o desempenho da rede em um cluster do Slurm
Monitorar VMs em um cluster do Slurm
Criar um endpoint de serviço: depois de ter o modelo refinado, implante-o em um endpoint de serviço usando o Google Kubernetes Engine (GKE) ou a Vertex AI para torná-lo acessível para inferência.
Saiba mais sobre os serviços de armazenamento recomendados para cargas de trabalho de IA e ML