Treinar o Qwen2 em um cluster Slurm A4

Neste tutorial, mostramos como treinar um modelo de linguagem grande (LLM) em um cluster do Slurm de vários nós e várias GPUs no Google Cloud. O modelo usado neste tutorial é baseado em um modelo de 1,5 bilhão de parâmetros do Qwen2. O cluster do Slurm usa duas máquinas virtuais (VMs) a4-highgpu-8g, 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 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, pesquisadores, administradores e operadores de plataforma de machine learning (ML), além de 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 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 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 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.

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.

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, 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:

    # Default configuration for a 2-node, 8-GPU-per-node (16 total GPUs) FSDP training job.

compute_environment: "LOCAL_MACHINE"
distributed_type: "FSDP"
downcast_bf16: "no"
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:

    #!/bin/bash
#SBATCH --job-name=qwen2-pretrain-smollm-fineweb
#SBATCH --nodes=2
#SBATCH --ntasks-per-node=8 # 8 tasks per node
#SBATCH --gpus-per-task=1   # 1 GPU per task
#SBATCH --partition=a4high
#SBATCH --output=logs/slurm-%j.out
#SBATCH --error=logs/slurm-%j.err

set -e
echo "--- Slurm Job Started ---"

# --- STAGE 1: Setup environment and pre-process data on each node's local SSD ---
# This command runs once per node.
srun --ntasks=$SLURM_NNODES --ntasks-per-node=1 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}"
  rsync -a --info=progress2 ~/./.venv/ ${LOCAL_VENV}/
  mkdir -p ${LOCAL_CACHE} ${PROCESSED_DATA_DIR}

  echo "Pre-processing data on $(hostname)..."
  source ${LOCAL_VENV}/bin/activate
  export HF_HOME=${LOCAL_CACHE}
  export HF_DATASETS_CACHE=${LOCAL_CACHE}

  # This runs the new preprocessing script. It ensures only ONE process per node
  # downloads and processes the data, avoiding rate limiting and redundant work.
  python 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}"
mkdir -p ${LOCAL_OUTPUT_DIR}

# This is the main training command. It launches one Python process per GPU.
srun --ntasks=$((SLURM_NNODES * 8)) --gpus-per-task=1 bash -c "
  source ${LOCAL_VENV}/bin/activate

  # The training script now loads the pre-processed data from the local SSD.
  python 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
"

# --- STAGE 3: Copy Final Model from Local SSD to Home Directory ---
echo "--- Copying final model from local SSD to /home ---"
# This command runs only on the first node of the job allocation
# and copies the final model back to the persistent shared directory.
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

  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 smollm-corpus dataset, e.g., 'fineweb-edu-dedup'.")
    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-olmo", 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)

    # --- 4. 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.")

    # --- 3. 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}...")
        lm_dataset = load_from_disk(args.preprocessed_data_path)
    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: sum(examples[k], []) 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. 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 \
  ./train.py \
  ./requirements.txt \
  ./submit.slurm \
  ./install_environment.sh \
  ./accelerate_config.yaml \
  "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. Crie uma variável de ambiente para o token de acesso do Hugging Face:

    export HUGGING_FACE_TOKEN="HUGGING_FACE_TOKEN"

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

    chmod +x install_environment.sh
./install_environment.sh

Começar a pré-treinar sua 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-qwen2-pretrain-smollm-fineweb.err

    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

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.

  • 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
LOGIN_NODE_NAME="your-login-node-name" # e.g., a4high-login-001
PROJECT_ID="your-gcp-project-id"
ZONE="your-cluster-zone" # e.g., us-west4-a

gcloud compute scp --project="$PROJECT_ID" --zone="$ZONE" --tunnel-through-iap \
  "${LOGIN_NODE_NAME}":~/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 do Hugging Face Transformers 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 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

A seguir