Treinar o Qwen2 em um cluster Slurm A4

Neste tutorial, você vai aprender a treinar um modelo de linguagem grande (LLM) em um cluster Slurm de vários nós e várias GPUs no Google Cloud. O modelo que você vai usar neste tutorial é baseado em um modelo Qwen2 de 1,5 bilhão de parâmetros. O cluster Slurm usa duas máquinas virtuais (VMs) a4-highgpu-8g, cada uma com 8 GPUs NVIDIA B200.

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

  1. Implante um cluster Slurm de alto desempenho 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 uma instância compartilhada do Filestore e uma rede RDMA de alta velocidade.
  2. Depois que o cluster for implantado, execute um job de pré-treinamento distribuído usando o conjunto de scripts que acompanham este tutorial. O job usa a biblioteca Accelerate do Hugging Face.

Este tutorial é destinado a engenheiros e pesquisadores de machine learning (ML), administradores e operadores de plataforma e especialistas em dados e IA interessados em implantar clusters Slurm de alta performance no Google Cloud para treinar LLMs.

Objetivos

  • Acesse o modelo Qwen2 usando o Hugging Face.
  • Prepare seu ambiente.
  • Crie e implante um cluster Slurm A4 de nível de produção.
  • Treine o modelo Qwen2 usando a biblioteca Accelerate .
  • 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 CLI gcloud com sua identidade federada.

  4. Para inicializar a CLI gcloud, 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 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

  8. Instale a CLI do Google Cloud.

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

  10. Para inicializar a CLI gcloud, 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 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
    14.

  14. 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, roles/compute.osAdminLogin, roles/iap.tunnelResourceAccessor 

    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.

Acessar o Qwen2 usando o Hugging Face

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

  1. Assine o contrato de consentimento para usar o Qwen 2 1.5B.

  2. Crie um token de acesso read.

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.

  1. 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

  2. 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"

  3. 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:

  1. Defina as variáveis de ambiente padrão: 

    export PROJECT_ID="YOUR_PROJECT_ID"
export CLUSTER_NAME="YOUR_CLUSTER_NAME"
export ZONE="YOUR_ZONE"
export REGION="YOUR_REGION"
export RESERVATION_URL="RESERVATION_NAME"
export GCS_BUCKET="YOUR_BUCKET_NAME"
export HF_TOKEN="HUGGING_FACE_TOKEN"

gcloud config set project "${PROJECT_ID}"
gcloud config set billing/quota_project "${PROJECT_ID}"

    Substitua:

    • YOUR_PROJECT_ID: o nome do Google Cloud projeto em que você quer criar o cluster do GKE.
    • YOUR_CLUSTER_NAME: o nome do cluster do Slurm que você quer criar.
    • YOUR_ZONE: a zona em que a reserva existe.
    • YOUR_REGION,: a região em que sua reserva existe.
    • RESERVATION_NAME: o URL ou o nome da reserva que você quer usar para criar o cluster do Slurm.
    • YOUR_BUCKET_NAME: o nome do bucket em que você armazena os resultados do checkpoint de treinamento. Antes de criar, familiarize-se com os requisitos de nomenclatura de bucket.
    • HUGGING_FACE_TOKEN: o token do Hugging Face que você criou em uma etapa anterior.

  2. Crie um bucket do Cloud Storage:

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

Criar um cluster Slurm A4

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

  1. Crie o arquivo a4high-slurm-deployment.yaml:

    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: ${GCS_BUCKET}

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

  2. Crie os manifestos do Terraform:

    gcluster create \
  -d "${MANIFEST_PATH}/a4high-slurm-deployment.yaml" \
  "${MANIFEST_PATH}/a4high-slurm-blueprint.yaml" \
  -o "${CLUSTER_NAME}"

  3. Adicione patches aos manifestos:

    sed -i '/deletion_protection = {/,/}/ { s/enabled = true/enabled = false/; /reason  = "Avoid data loss"/d; }' "${CLUSTER_NAME}/${CLUSTER_NAME}/cluster-env/main.tf"

  4. Implante o cluster:

    gcluster deploy "${CLUSTER_NAME}/${CLUSTER_NAME}" --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é 50 minutos.

    • A segunda fase implanta o cluster usando essa imagem personalizada. Esse processo geralmente leva menos tempo para ser concluído 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 "${CLUSTER_NAME}" --auto-approve --skip "image" -w

Preparar sua carga de trabalho

Para preparar sua carga de trabalho, siga estas etapas:

  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 treinamento 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 should be run ONCE on the login node to set up the
# shared Python virtual environment.

set -e
echo "--- Creating Python virtual environment in /home ---"
python3 -m venv ~/.venv
echo "--- Activating virtual environment ---"
source ~/.venv/bin/activate

echo "--- Installing build dependencies ---"
pip install --upgrade pip wheel packaging

echo "--- Installing PyTorch for CUDA 12.8 ---"
pip install torch --index-url https://download.pytorch.org/whl/cu128

echo "--- Installing application requirements ---"
pip install -r requirements.txt

echo "--- Environment setup complete. You can now submit jobs with sbatch. ---"

  2. Para especificar as configurações do job de ajuste refinado, crie o arquivo accelerate_config.yaml com o seguinte conteúdo:

    compute_environment: "LOCAL_MACHINE"
distributed_type: "FSDP"
downcast_bf16: "no"
fsdp_config:
  fsdp_auto_wrap_policy: "TRANSFORMER_BASED_WRAP"
  fsdp_backward_prefetch: "BACKWARD_PRE"
  fsdp_cpu_ram_efficient_loading: true
  fsdp_forward_prefetch: false
  fsdp_offload_params: false
  fsdp_sharding_strategy: "FULL_SHARD"
  fsdp_state_dict_type: "SHARDED_STATE_DICT"
  fsdp_transformer_layer_cls_to_wrap: "Qwen2DecoderLayer"
machine_rank: 0
main_training_function: "main"
mixed_precision: "bf16"
num_machines: 2
num_processes: 16
rdzv_backend: "static"
same_network: true
tpu_env: []
use_cpu: false

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

    #SBATCH --job-name=qwen2-pretrain-smollm-fineweb
#SBATCH --nodes=2
#SBATCH --ntasks-per-node=1  # 1 main srun task on the node that manages accelerate
#SBATCH --gpus-per-node=8    # access to all 8 GPUs on the node
#SBATCH --partition=a4high
#SBATCH --output=logs/slurm-%j.out
#SBATCH --error=logs/slurm-%j.err

set -euo pipefail
echo "--- Slurm Job Started ---"

# --- STAGE 1: Setup environment and pre-process data on each node's local SSD ---
srun --ntasks=$SLURM_NNODES --ntasks-per-node=1 --gpu-bind=none bash -c '
  set -e
  echo "Setting up local environment on $(hostname)..."
  LOCAL_VENV="/mnt/localssd/venv_job_${SLURM_JOB_ID}"
  LOCAL_CACHE="/mnt/localssd/hf_cache_job_${SLURM_JOB_ID}"
  PROCESSED_DATA_DIR="/mnt/localssd/processed_data_${SLURM_JOB_ID}"
  LOCAL_TMP="/mnt/localssd/tmp_job_${SLURM_JOB_ID}"

  rsync -a --info=progress2 ~/.venv/ ${LOCAL_VENV}/
  mkdir -p ${LOCAL_CACHE} ${PROCESSED_DATA_DIR} ${LOCAL_TMP}

  echo "Pre-processing data on $(hostname)..."
  source ${LOCAL_VENV}/bin/activate
  export TMPDIR="${LOCAL_TMP}"
  export TEMP="${LOCAL_TMP}"
  export TMP="${LOCAL_TMP}"
  export HF_HOME="${LOCAL_CACHE}"
  export HF_DATASETS_CACHE="${LOCAL_CACHE}"
  export HF_MODULES_CACHE="${LOCAL_CACHE}/modules"
  export HF_METRICS_CACHE="${LOCAL_CACHE}/metrics"

  python "${HOME}/preprocess_data.py" \
    --dataset_name "HuggingFaceFW/fineweb-edu" \
    --dataset_config "CC-MAIN-2024-10" \
    --tokenizer_id "Qwen/Qwen2-1.5B" \
    --max_seq_length 1024 \
    --output_path ${PROCESSED_DATA_DIR}

  echo "Setup on $(hostname) complete."
'

# --- STAGE 2: Run the Training Job using the Local Environment ---
echo "--- Starting Training ---"

LOCAL_VENV="/mnt/localssd/venv_job_${SLURM_JOB_ID}"
PROCESSED_DATA_DIR="/mnt/localssd/processed_data_${SLURM_JOB_ID}"
LOCAL_OUTPUT_DIR="/mnt/localssd/outputs_${SLURM_JOB_ID}"

export MASTER_ADDR=$(scontrol show hostnames "$SLURM_JOB_NODELIST" | head -n 1)
export MASTER_PORT=29505

# Network and initialization debugging configurations
export NCCL_DEBUG=INFO
export NCCL_DEBUG_SUBSYS=INIT,COLL
export TORCH_DISTRIBUTED_DEBUG=INFO
export NCCL_IB_DISABLE=0

# Launching with full GPU access enabled for accelerate
srun --ntasks=$SLURM_NNODES --ntasks-per-node=1 --gpu-bind=none bash -c "
  mkdir -p ${LOCAL_OUTPUT_DIR}
  source ${LOCAL_VENV}/bin/activate

  # Reset the forced Slurm isolation to expose all 8 GPUs to the processes
  unset CUDA_VISIBLE_DEVICES

  # Retrieve the default network device.
  DETECTED_IFACE=\$(ip route show | grep default | awk '{print \$5}' | head -n 1)
  echo \"[INFO] Automatically detected cluster network interface: \$DETECTED_IFACE\"

  # Dynamic injection of the detected interface into the network configuration
  export NCCL_SOCKET_IFNAME=\"\${DETECTED_IFACE},gpu*\"
  export TP_SOCKET_IFNAME=\"\${DETECTED_IFACE}\"
  export GLOO_SOCKET_IFNAME=\"\${DETECTED_IFACE}\"

  accelerate launch \
    --config_file ~/accelerate_config.yaml \
    --num_machines \$SLURM_NNODES \
    --num_processes \$((SLURM_NNODES * 8)) \
    --machine_rank \$SLURM_NODEID \
    --main_process_ip \$MASTER_ADDR \
    --main_process_port \$MASTER_PORT \
    train.py \
      --model_config_id 'Qwen/Qwen2-1.5B' \
      --preprocessed_data_path ${PROCESSED_DATA_DIR} \
      --output_dir ${LOCAL_OUTPUT_DIR} \
      --per_device_train_batch_size 4 \
      --gradient_accumulation_steps 4 \
      --max_steps 10000 \
      --learning_rate 5e-5 \
      --save_strategy steps \
      --save_steps 500 \
      --logging_steps 1
"

# --- STAGE 3: Copy Final Model from Local SSD to Home Directory ---
echo "--- Copying final model from local SSD to /home ---"
mkdir -p ~/qwen2-from-scratch-on-smollm-fineweb/

srun --nodes=1 --ntasks=1 --ntasks-per-node=1 bash -c "
  rsync -a --info=progress2 ${LOCAL_OUTPUT_DIR}/ ~/qwen2-from-scratch-on-smollm-fineweb/
"

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

  4. Para especificar as dependências do job de ajuste refinado, crie um arquivo requirements.txt com o seguinte conteúdo:

    # Hugging Face Libraries (Pinned to recent, stable versions for reproducibility)
transformers==4.53.3
datasets==4.0.0
accelerate==1.9.0
evaluate==0.4.5
bitsandbytes==0.46.1
trl==0.19.1
peft==0.16.0

# Other dependencies
tensorboard==2.20.0
protobuf==6.31.1
sentencepiece==0.2.0
    Observação: este tutorial não usa as versões mais recentes das dependências do NVIDIA e do PyTorch. Se você precisar de dependências mais recentes, consulte a documentação da NVIDIA e a documentação do Pytorch.

  5. Para fazer o download, tokenizar e pré-processar o conjunto de dados em um formato pronto para treinamento, crie um arquivo preprocess_data.py com o seguinte conteúdo:

    import argparse
from datasets import load_dataset
from transformers import AutoTokenizer
import os
from itertools import chain

def get_args():
   parser = argparse.ArgumentParser(description="Download and preprocess a dataset.")
   parser.add_argument("--dataset_name", type=str, required=True)
   parser.add_argument("--dataset_config", type=str, required=True)
   parser.add_argument("--tokenizer_id", type=str, required=True)
   parser.add_argument("--max_seq_length", type=int, required=True)
   parser.add_argument("--output_path", type=str, required=True, help="Path to save the processed dataset.")
   return parser.parse_args()

def main():
   args = get_args()

   if os.path.exists(args.output_path) and os.listdir(args.output_path):
       print(f"Processed dataset already exists at {args.output_path}. Skipping.")
       return

   # 1. Load tokenizer
   tokenizer = AutoTokenizer.from_pretrained(args.tokenizer_id)

   # 2. Load raw dataset
   print(f"Loading raw dataset {args.dataset_name}...")
   raw_dataset = load_dataset(args.dataset_name, name=args.dataset_config, split="train")

   # 3. Tokenize
   def tokenize_function(examples):
       return tokenizer(examples["text"])

   num_proc = os.cpu_count()
   print(f"Tokenizing dataset using {num_proc} processes...")
   print("Tokenizing dataset...")
   tokenized_dataset = raw_dataset.map(
       tokenize_function,
       batched=True,
       remove_columns=raw_dataset.column_names,
       desc="Running tokenizer on dataset",
       num_proc=num_proc,
   )

   # 4. Group texts
   def group_texts(examples):
       concatenated_examples = {k: list(chain.from_iterable(examples[k])) for k in examples.keys()}
       total_length = len(concatenated_examples[list(examples.keys())[0]])
       if total_length >= args.max_seq_length:
           total_length = (total_length // args.max_seq_length) * args.max_seq_length
       result = {
           k: [t[i : i + args.max_seq_length] for i in range(0, total_length, args.max_seq_length)]
           for k, t in concatenated_examples.items()
       }
       result["labels"] = result["input_ids"].copy()
       return result

   print("Grouping texts...")
   lm_dataset = tokenized_dataset.map(
       group_texts,
       batched=True,
       desc=f"Grouping texts in chunks of {args.max_seq_length}",
       num_proc=num_proc,
   )

   # 5. Save to disk
   print(f"Saving processed dataset to {args.output_path}...")
   lm_dataset.save_to_disk(args.output_path)
   print("Preprocessing complete.")

if __name__ == "__main__":
   main()

  6. Para especificar as instruções do job, crie um arquivo train.py com o seguinte conteúdo:

    import torch
import argparse
from datasets import load_dataset, load_from_disk
import os
from transformers import (
    AutoConfig,
    AutoTokenizer,
    AutoModelForCausalLM,
    Trainer,
    TrainingArguments,
    DataCollatorForLanguageModeling,
)
from huggingface_hub import login

def get_args():
    parser = argparse.ArgumentParser()
    parser.add_argument("--model_config_id", type=str, default="Qwen/Qwen2-1.5B", help="Hugging Face model config to use for architecture.")
    # Data arguments - used if preprocessed data is not available
    parser.add_argument("--dataset_name", type=str, default="HuggingFaceFW/fineweb-edu", help="Hugging Face dataset for pre-training.")
    parser.add_argument("--dataset_config", type=str, default="CC-MAIN-2024-10", help="Config for the fineweb-edu dataset, e.g., 'CC-MAIN-2024-10'.")
    parser.add_argument("--preprocessed_data_path", type=str, default=None, help="Path to a preprocessed dataset on disk. If provided, skips download and processing.")
    # General arguments
    parser.add_argument("--hf_token", type=str, default=None, help="Hugging Face token for private models/tokenizers")
    parser.add_argument("--output_dir", type=str, default="qwen2-from-scratch-on-smollm-fineweb", help="Directory to save model checkpoints")

    # TrainingArguments
    parser.add_argument("--max_seq_length", type=int, default=1024, help="Maximum sequence length")
    parser.add_argument("--num_train_epochs", type=int, default=1, help="Number of training epochs")
    parser.add_argument("--max_steps", type=int, default=-1, help="If set to a positive number, it overrides num_train_epochs.")
    parser.add_argument("--per_device_train_batch_size", type=int, default=4, help="Batch size per device during training")
    parser.add_argument("--gradient_accumulation_steps", type=int, default=4, help="Gradient accumulation steps")
    parser.add_argument("--learning_rate", type=float, default=5e-5, help="Learning rate")
    parser.add_argument("--logging_steps", type=int, default=10, help="Log every X steps")
    parser.add_argument("--save_strategy", type=str, default="steps", help="Checkpoint save strategy")
    parser.add_argument("--save_steps", type=int, default=500, help="Save checkpoint every X steps")

    return parser.parse_args()

def main():
    args = get_args()

    # --- 1. Setup and Login ---
    if args.hf_token:
        login(args.hf_token)

    # --- 2. Load Tokenizer ---
    # We load the tokenizer from the specified config ID to ensure compatibility
    # with the model architecture (e.g., special tokens).
    tokenizer = AutoTokenizer.from_pretrained(args.model_config_id)

    # --- 3. Initialize Model from Scratch ---
    print(f"Initializing a new model from {args.model_config_id} configuration...")
    config = AutoConfig.from_pretrained(args.model_config_id)
    model = AutoModelForCausalLM.from_config(config)

    print(f"Model has {model.num_parameters():,} parameters.")

    # --- 4. Load or Create and prepare the training dataset ---
    if args.preprocessed_data_path and os.path.exists(args.preprocessed_data_path):
        print(f"Loading preprocessed dataset from {args.preprocessed_data_path}...")

        # Synchronization of distributed processes
        local_rank = int(os.environ.get("LOCAL_RANK", -1))
        if local_rank != -1:
            # Introducing a minimal time offset per GPU to avoid I/O collisions.
            import time
            time.sleep(local_rank * 0.2)

        lm_dataset = load_from_disk(args.preprocessed_data_path, keep_in_memory=False)

    else:
        print("No preprocessed dataset found, starting from raw data...")
        raw_dataset = load_dataset(args.dataset_name, name=args.dataset_config, split="train")

        # Tokenization function
        def tokenize_function(examples):
            return tokenizer(examples["text"])

        tokenized_dataset = raw_dataset.map(
            tokenize_function,
            batched=True,
            remove_columns=raw_dataset.column_names,
            desc="Running tokenizer on dataset",
        )

        # Main data processing function that will concatenate all texts from our dataset
        # and generate chunks of max_seq_length.
        def group_texts(examples):
            # Concatenate all texts.
            concatenated_examples = {k: [item for sublist in examples[k] for item in sublist] for k in examples.keys()}
            total_length = len(concatenated_examples[list(examples.keys())[0]])
            # We drop the small remainder.
            if total_length >= args.max_seq_length:
                total_length = (total_length // args.max_seq_length) * args.max_seq_length
            # Split by chunks of max_len.
            result = {
                k: [t[i : i + args.max_seq_length] for i in range(0, total_length, args.max_seq_length)]
                for k, t in concatenated_examples.items()
            }
            result["labels"] = result["input_ids"].copy()
            return result

        lm_dataset = tokenized_dataset.map(
            group_texts,
            batched=True,
            desc=f"Grouping texts in chunks of {args.max_seq_length}",
        )


    # --- 5. Configure Training Arguments ---
    # Check for bfloat16 support
    use_bf16 = torch.cuda.is_available() and torch.cuda.is_bf16_supported()

    training_args = TrainingArguments(
        output_dir=args.output_dir,
        num_train_epochs=args.num_train_epochs,
        max_steps=args.max_steps,
        per_device_train_batch_size=args.per_device_train_batch_size,
        gradient_accumulation_steps=args.gradient_accumulation_steps,
        learning_rate=args.learning_rate,
        logging_steps=args.logging_steps,
        save_strategy=args.save_strategy,
        save_steps=args.save_steps,
        save_total_limit=2, # Optional: Limit the number of checkpoints
        bf16=use_bf16,
        fp16=not use_bf16,
        optim="adamw_torch",
        lr_scheduler_type="cosine",
        warmup_ratio=0.03,
        report_to="tensorboard",
        gradient_checkpointing=True,
        # Required for gradient checkpointing with some parallelization strategies
        gradient_checkpointing_kwargs={"use_reentrant": False},
    )

    # --- 6. Create Trainer and Start Training ---
    # Data collator will take care of creating batches for causal language modeling
    data_collator = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False)

    trainer = Trainer(
        model=model,
        args=training_args,
        train_dataset=lm_dataset,
        # eval_dataset=... # Optional: if you have a validation set
        tokenizer=tokenizer,
        data_collator=data_collator,
    )

    print("Starting training from scratch...")
    trainer.train()
    print("Training finished.")

    # --- 7. Save the final model ---
    print(f"Saving final model to {args.output_dir}")
    trainer.save_model()

if __name__ == "__main__":
    main()

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. 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='${CLUSTER_NAME}' AND labels.slurm_instance_role='login'" \
                --format="value(name)" | head -n 1)"

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

  2. 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)"

  3. 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.txt \
  ./submit.slurm \
  ./accelerate_config.yaml \
  ./train.py \
  ./preprocess_data.py \
  "${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 frameworks e ferramentas fazendo o seguinte:

  1. Configure um ambiente virtual do Python com todas as dependências necessárias:

    chmod +x install_environment.sh
./install_environment.sh

Começar o pré-treinamento da carga de trabalho

Para iniciar o treinamento da 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:

    tail -f ~/logs/slurm-1.err # (or .out, depending on where the script is currently sending logs)

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

Monitore sua carga de trabalho

É 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=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 de GPUs: 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. TESTAR

  • Duração do job: o job leva aproximadamente uma hora para ser concluído.

Baixar o modelo

Depois que o job for executado, o modelo treinado será salvo no diretório ~/qwen2-from-scratch-on-smollm-fineweb/ no nó de login. Como esse diretório compartilhado persistente é montado em todos os nós do cluster, os pontos de verificação do modelo permanecem disponíveis mesmo depois que o job é concluído ou que os nós de computação são desalocados.

É possível fazer o download do modelo salvo do nó de login para sua máquina local usando o comando gcloud compute scp, conforme mostrado no exemplo a seguir:

# From your local machine
gcloud compute scp --project="${PROJECT_ID}" --zone="${ZONE}" --tunnel-through-iap \
  "${LOGIN_NODE}":~/qwen2-from-scratch-on-smollm-fineweb/ ./qwen2-trained-model/ --recurse

Depois de baixar o modelo, você pode fazer o seguinte:

  • Carregue o modelo para inferência: use a estrutura Transformers do Hugging Face para carregar o diretório qwen2-trained-model/ e realizar a inferência com seu modelo Qwen2 recém-treinado.
  • Ajuste refinado adicional: use o ponto de verificação salvo como ponto de partida para um ajuste refinado adicional em um conjunto de dados mais específico.
  • Envie o modelo para o Hugging Face Hub: compartilhe seu modelo treinado enviando-o para o Hugging Face Hub.

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

  1. Para excluir o cluster do Slurm, siga estas etapas:

    ./gcluster destroy "${CLUSTER_NAME}" --auto-approve

  2. Para excluir o bucket do Cloud Storage, siga estas etapas:

    gcloud storage buckets delete "gs://${GCS_BUCKET}" --quiet || true

  3. Para excluir uma imagem do Packer, abra o navegador da Web, navegue até a página a seguir, pesquise sua imagem específica e clique em "Excluir". 

    http://console.cloud.google.com/compute/images

  4. Para excluir todas as redes VPC, regras de firewall, roteadores, IPs e sub-redes associados ao projeto, faça o seguinte:

    echo "========================================================================="
echo " STARTING AUTOMATED NETWORK CLEANUP FOR CLUSTER: ${CLUSTER_NAME}"
echo "========================================================================="

echo "Discovering all VPC networks linked to the cluster..."
NETWORKS=$(gcloud compute networks list --project="${PROJECT_ID}" --format="value(name)" | grep "^${CLUSTER_NAME}" || true)

if [ -z "${NETWORKS}" ]; then
    echo "No VPC networks found starting with ${CLUSTER_NAME}. Everything is already clean!"
    exit 0
fi

echo "Found the following networks to process:"
echo "${NETWORKS}"
echo "-------------------------------------------------------------------------"

echo "=== 1. Wiping Global Firewall Rules ==="
FIREWALL_RULES=$(gcloud compute firewall-rules list \
    --project="${PROJECT_ID}" \
    --filter="network ~ ^${CLUSTER_NAME} OR name ~ ^${CLUSTER_NAME}" \
    --format="value(name)" || echo "")

if [ -n "${FIREWALL_RULES}" ]; then
    echo "Deleting matching firewall rules:"
    echo "${FIREWALL_RULES}"
    echo "${FIREWALL_RULES}" | xargs -r gcloud compute firewall-rules delete --project="${PROJECT_ID}" --quiet
else
    echo "No matching firewall rules found."
fi

echo "=== 2. Tearing Down Network-Specific Infrastructure ==="
echo "${NETWORKS}" | while read -r net_name; do
    [ -z "${net_name}" ] && continue
    echo "Processing resources for network: ${net_name}"

    ROUTERS=$(gcloud compute routers list \
        --project="${PROJECT_ID}" \
        --regions="${REGION}" \
        --filter="network=${net_name}" \
        --format="value(name)" || echo "")

    if [ -n "${ROUTERS}" ]; then
        echo "  -> Deleting routers: ${ROUTERS}"
        echo "${ROUTERS}" | xargs -r gcloud compute routers delete --region="${REGION}" --project="${PROJECT_ID}" --quiet
    fi

    IPS=$(gcloud compute addresses list \
        --project="${PROJECT_ID}" \
        --regions="${REGION}" \
        --filter="name ~ ^${net_name}" \
        --format="value(name)" || echo "")

    if [ -n "${IPS}" ]; then
        echo "  -> Deleting IP reservations: ${IPS}"
        echo "${IPS}" | xargs -r gcloud compute addresses delete --region="${REGION}" --project="${PROJECT_ID}" --quiet
    fi

    SUBNETS=$(gcloud compute networks subnets list \
        --project="${PROJECT_ID}" \
        --regions="${REGION}" \
        --filter="network=${net_name}" \
        --format="value(name)" || echo "")

    if [ -n "${SUBNETS}" ]; then
        echo "  -> Deleting subnetworks:"
        echo "${SUBNETS}"
        echo "${SUBNETS}" | xargs -r gcloud compute networks subnets delete --region="${REGION}" --project="${PROJECT_ID}" --quiet
    fi
done

echo "-------------------------------------------------------------------------"
echo "Waiting 15 seconds for Google Cloud API dependencies to unlock..."
sleep 15

echo "=== 3. Final VPC Networks Destruction ==="
echo "${NETWORKS}" | while read -r net_name; do
    [ -z "${net_name}" ] && continue
    echo "Deleting core VPC network: ${net_name}..."
    gcloud compute networks delete "${net_name}" --project="${PROJECT_ID}" --quiet || \
    echo "Warning: Could not delete ${net_name} yet. If a lock occurred, please rerun in 1 minute."
done

echo "========================================================================="
echo " SUCCESS: All network resources for cluster ${CLUSTER_NAME} have been wiped!"
echo "========================================================================="

Excluir o projeto

Excluir um projeto do Google Cloud :

gcloud projects delete PROJECT_ID

A seguir