Ajuste fino do modelo Olmo3 7B com Axolotl em infraestrutura sem servidor e com múltiplos GPUs

Abrir a versão desta página do bloco de notas

Afina finamente o modelo Olmo3 7B Instruct no tempo de execução da IA usando o Axolotl. O Axolotl fornece uma framework de alto desempenho para o pós-treino de LLM com QLoRA (Quantized Low-Rank Adaptation), permitindo um ajuste fino eficiente em infraestruturas multi-GPU. O modelo treinado é registado no MLflow e registado no Unity Catalog para implementação.

Ligar-se à computação GPU sem servidor

Este notebook requer computação GPU sem servidor. Para ligar:

  1. Clique no seletor de computação do portátil no canto superior direito e selecione GPU Serverless
  2. No lado direito, clique no botão de ambiente
  3. Selecione o 8xH100 como acelerador
  4. Selecione AI v5 como ambiente e depois clique em Aplicar

Instalar dependências necessárias

Instala o Axolotl com suporte para Flash Attention e versões compatíveis da trl e de bibliotecas de otimização. O cut-cross-entropy pacote fornece cálculo de perda eficiente em memória para grandes modelos de linguagem.

%pip install --no-build-isolation "axolotl[flash-attn]==0.13.1"
%pip install "trl==0.27.1"
%pip install "torchao==0.16.0"
%pip install "cut-cross-entropy[transformers] @ git+https://github.com/axolotl-ai-cloud/ml-cross-entropy.git@f4b5712"
dbutils.library.restartPython()

Recuperar o token HuggingFace

Recupera o token de autenticação da HuggingFace dos segredos do Databricks. Este token é necessário para descarregar o modelo base Olmo3 7B do HuggingFace Hub.

HF_TOKEN = dbutils.secrets.get(scope="sgc-nightly-notebook", key="hf_token")

Configurar parâmetros de treino

Configura a configuração de treino do Axolote com base no exemplo olmo3-7b-qlora.yaml . As principais modificações incluem:

  • Integração com MLflow para rastreamento de experiências
  • Caminho de volume do Unity Catalog para armazenamento de pontos de verificação
  • SDPA (Scaled Dot Product Attention) em vez de Flash Attention para uma compatibilidade mais ampla com GPUs

Definir caminhos do Catálogo Unity

Cria widgets para especificar a localização do Unity Catalog para o armazenamento de pontos de verificação do modelo. O diretório de saída combina o catálogo, o esquema, o volume e o nome do modelo num caminho totalmente qualificado.

dbutils.widgets.text("uc_catalog", "main")
dbutils.widgets.text("uc_schema", "default")
dbutils.widgets.text("uc_volume", "checkpoints")
dbutils.widgets.text("model", "openai/gpt-oss-20b")

UC_CATALOG = dbutils.widgets.get("uc_catalog")
UC_SCHEMA = dbutils.widgets.get("uc_schema")
UC_VOLUME = dbutils.widgets.get("uc_volume")
UC_MODEL_NAME = dbutils.widgets.get("model")

print(f"UC_CATALOG: {UC_CATALOG}")
print(f"UC_SCHEMA: {UC_SCHEMA}")
print(f"UC_VOLUME: {UC_VOLUME}")
print(f"UC_MODEL_NAME: {UC_MODEL_NAME}")

OUTPUT_DIR = f"/Volumes/{UC_CATALOG}/{UC_SCHEMA}/{UC_VOLUME}/{UC_MODEL_NAME}"
print(f"OUTPUT_DIR: {OUTPUT_DIR}")

Desativar telemetria

Desativa o rastreamento de uso do Axolotl ao definir a variável de ambiente.

import os
os.environ['AXOLOTL_DO_NOT_TRACK'] = '1'

Criar configuração Axolotl

Define a configuração completa de treino usando o formato DictDefault do Axolotl. Isto inclui definições de modelo (QLoRA com quantização de 4 bits), configuração do conjunto de dados (formato Alpaca), hiperparâmetros LoRA (rank 32, alpha 16), parâmetros de treino (1 época, tamanho do lote 2, acumulação de gradiente 4) e integração MLflow para rastreamento de experiências.

from axolotl.cli.config import load_cfg
from axolotl.utils.dict import DictDefault

# Config is based on with some changes to fit GPU types
# https://raw.githubusercontent.com/axolotl-ai-cloud/axolotl/main/examples/olmo3/olmo3-7b-qlora.yaml

# Axolotl provides full control and transparency over model and training configuration
config = DictDefault(
    base_model="allenai/Olmo-3-7B-Instruct-SFT",
    plugins=[
        "axolotl.integrations.cut_cross_entropy.CutCrossEntropyPlugin"
    ],
    load_in_8bit=False,
    load_in_4bit=True,
    datasets=[
        {
            "path": "fozziethebeat/alpaca_messages_2k_test",
            "type": "chat_template"
        }
    ],
    dataset_prepared_path="last_run_prepared",
    val_set_size=0.1,
    output_dir=OUTPUT_DIR,
    adapter="qlora",
    lora_model_dir=None,
    sequence_len=2048,
    sample_packing=True,
    lora_r=32,
    lora_alpha=16,
    lora_dropout=0.05,
    lora_target_linear=True,
    lora_target_modules=[
        "gate_proj",
        "down_proj",
        "up_proj",
        "q_proj",
        "v_proj",
        "k_proj",
        "o_proj"
    ],
    wandb_project=None,
    wandb_entity=None,
    wandb_watch=None,
    wandb_name=None,
    wandb_log_model=None,
    gradient_accumulation_steps=4,
    micro_batch_size=2,
    num_epochs=1,
    optimizer="adamw_bnb_8bit",
    lr_scheduler="cosine",
    learning_rate=0.0002,
    bf16="auto",
    tf32=False,
    gradient_checkpointing=True,
    resume_from_checkpoint=None,
    logging_steps=1,
    flash_attention=False,
    warmup_ratio=0.1,
    evals_per_epoch=1,
    saves_per_epoch=1,
    # Eval dataset is too small
    eval_sample_packing=False,
    # Write metrics to MLflow
    use_mlflow=True,
    mlflow_tracking_uri="databricks",
    mlflow_run_name="olmo3-7b-qlora-axolotl",
    hf_mlflow_log_artifacts=False,
    wandb_mode="disabled",
    attn_implementation="sdpa",
    sdpa_attention=True,
    save_first_step=True,
    device_map=None,
)

Configurar a alocação de memória CUDA do PyTorch

Otimiza a gestão da memória GPU para treino eficiente em configurações multi-GPU.

from axolotl.utils import set_pytorch_cuda_alloc_conf

set_pytorch_cuda_alloc_conf()

Executar treino distribuído em computação com GPU serverless

Usa o @distributed decorador da API de GPU serverless para distribuir a tarefa de treino do Axolotl por 8 GPUs H100. O decorador lida com a orquestração de várias GPUs, permitindo que a função de treino funcione num ambiente distribuído sem a configuração manual do cluster.

from serverless_gpu.launcher import distributed
from serverless_gpu.compute import GPUType

@distributed(gpus=8, gpu_type=GPUType.H100)
def run_train(cfg: DictDefault):
    import os
    os.environ['HF_TOKEN'] = HF_TOKEN

    from axolotl.common.datasets import load_datasets

    # Load, parse and tokenize the datasets to be formatted with qwen3 chat template
    # Drop long samples from the dataset that overflow the max sequence length

    # validates the configuration
    cfg = load_cfg(cfg)
    dataset_meta = load_datasets(cfg=cfg)

    from axolotl.train import train

    # just train the first 16 steps for demo.
    # This is sufficient to align the model as we've used packing to maximize the trainable samples per step.
    cfg.max_steps = 16
    model, tokenizer, trainer = train(cfg=cfg, dataset_meta=dataset_meta)

    import mlflow
    mlflow_run_id = None
    if mlflow.last_active_run() is not None:
        mlflow_run_id = mlflow.last_active_run().info.run_id

    return mlflow_run_id

result = run_train.distributed(config)

Executar a tarefa de formação

Lança o trabalho de formação distribuída. A função carrega o conjunto de dados, valida a configuração, treina o modelo durante 16 passos e devolve o ID da corrida MLflow para rastreio.

run_id = result[0]
print(run_id)

Extrair o ID da execução MLflow

Recupera o run ID do MLflow dos resultados do treino para registo de modelos e monitorização de experiências.

Registe o modelo finamente ajustado no Unity Catalog

Carrega o adaptador LoRA treinado, funde-o com o modelo base e regista o modelo combinado no Unity Catalog via MLflow. Isto torna o modelo disponível para implementação e inferência.

Nota: Esta etapa requer o cálculo da GPU H100 para carregar o checkpoint do modelo. Correr em GPUs mais pequenas pode resultar em erros de CUDA por falta de memória.

from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline

from peft import PeftModel
import mlflow
import torch

HF_MODEL_NAME = "allenai/Olmo-3-7B-Instruct-SFT"

torch.cuda.empty_cache()
# Load the trained model for registration
print("Loading LoRA model for registration...")
# For LoRA models, we need both base model and adapter
base_model = AutoModelForCausalLM.from_pretrained(
    HF_MODEL_NAME,
    trust_remote_code=True
)
# Load tokenizer
tokenizer = AutoTokenizer.from_pretrained(HF_MODEL_NAME)
adapter_dir = OUTPUT_DIR
peft_model = PeftModel.from_pretrained(base_model, adapter_dir)
# Merge LoRA into base and drop PEFT wrappers
merged_model = peft_model.merge_and_unload()
merged_model.generation_config.temperature = None
merged_model.generation_config.top_p = None

# Create Unity Catalog model name
full_model_name = f"{UC_CATALOG}.{UC_SCHEMA}.{UC_MODEL_NAME}"

print(f"Registering model as: {full_model_name}")

text_gen_pipe = pipeline(
    task="text-generation",
    model=merged_model,
    tokenizer=tokenizer,
)

input_example = ["Hello, world!"]

with mlflow.start_run(run_id=run_id):
    model_info = mlflow.transformers.log_model(
        transformers_model=text_gen_pipe,
        name="model",
        input_example=input_example,
        registered_model_name=full_model_name,
    )
print(f"✓ Model successfully registered in Unity Catalog: {full_model_name}")
print(f"✓ MLflow model URI: {model_info.model_uri}")
print(f"✓ Model version: {model_info.registered_model_version}")

print(f"\n📦 Model Registration Complete!")
print(f"Unity Catalog Path: {full_model_name}")
print(f"Optimization: Cut Cross Entropy + QLoRA")

Passos seguintes

Exemplo de bloco de notas

Ajuste fino do modelo Olmo3 7B com Axolotl em infraestrutura sem servidor e com múltiplos GPUs

Obter bloco de notas

  • Last updated on