Pular para o conteúdo principal
Última atualização em

Ajuste fino do Olmo3 7B com Axolotl em compute serverless múltiplas GPUs.

Este notebook demonstra como ajustar o modelo Olmo3 7B Instruct usando o Axolotl em compute GPU serverless Databricks . O Axolotl fornece uma estrutura de alto desempenho para pós-treinamento de modelos de aprendizado de máquina (LLM) com QLoRA (Quantized Low-Rank Adaptation), permitindo um ajuste fino eficiente em infraestruturas com múltiplas GPUs. O modelo treinado é enviado para MLflow e registrado no Unity Catalog para implantação.

Instale as dependências necessárias​

Instala o Axolotl com suporte a Flash Attention, MLflow para acompanhamento de experimentos e versões compatíveis de transformadores e biblioteca de otimização. O pacote cut-cross-entropy fornece computação de perda com uso eficiente de memória para grandes modelos de linguagem.

Python
%pip install -U packaging setuptools wheel ninja
%pip install mlflow>=3.6
%pip install --no-build-isolation axolotl[flash-attn]>=0.12.0
%pip install transformers==4.57.3
%pip uninstall -y awq autoawq
%pip install "cut-cross-entropy[transformers] @ git+https://github.com/axolotl-ai-cloud/ml-cross-entropy.git@f4b5712"
dbutils.library.restartPython()

Recuperar tokens HuggingFace​

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

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

Configurar parâmetros de treinamento​

Configura o ambiente de treinamento do Axolotl com base no exemplo olmo3-7b-qlora.yaml . As principais modificações incluem:

  • Integração MLflow para acompanhamento de experimentos
  • Caminho do volume Unity Catalog para armazenamento de pontos de verificação.
  • SDPA (Scaled Dot produto Attention) em vez de Flash Attention para maior compatibilidade com GPUs.

Definir caminhos Unity Catalog​

Cria widgets para especificar a localização Unity Catalog para armazenar os 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 em um caminho totalmente qualificado.

Python
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}")
Python
import os
os.environ['AXOLOTL_DO_NOT_TRACK'] = '1'

Desativar telemetria​

Desabilita o acompanhamento de uso do Axolotl definindo a variável de ambiente.

Criar configuração do Axolotl​

Define a configuração completa de treinamento usando o formato DictDefault do Axolotl. Isso inclui configurações do modelo (QLoRA com quantização de 4 bits), configuração dataset (formato Alpaca), hiperparâmetros LoRA (rank 32, alpha 16), parâmetros de treinamento (1 época, tamanho do lote 2, acumulação de gradiente 4) e integração MLflow para acompanhamento do experimento.

Python
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,
)
Python
from axolotl.utils import set_pytorch_cuda_alloc_conf

set_pytorch_cuda_alloc_conf()

Configurar alocação de memória CUDA para PyTorch​

Otimiza o gerenciamento de memória da GPU para um treinamento eficiente em configurações com múltiplas GPUs.

execução distribuída treinamento em computeGPU serverless​

Utiliza o decorador @distributed da API de GPU serverless para distribuir o Job de treinamento do Axolotl entre 8 GPUs H100. O decorador lida com orquestração multi-GPU, permitindo que a função de treinamento seja executada em um ambiente distribuído sem configuração manual cluster .

Python
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
Python
result = run_train.distributed(config)

Executar o trabalho de treinamento​

Inicia a tarefa de treinamento distribuído. A função carrega o dataset, valida a configuração, ensina o modelo para 16 passos e retorna o ID de execução MLflow para acompanhamento.

Python
run_id = result[0]
print(run_id)

Extrair o ID da execução do MLflow​

Recupera o ID de execução MLflow dos resultados do treinamento para registro do modelo e acompanhamento do experimento.

Registre o modelo otimizado no Unity Catalog​

Carrega o adaptador LoRa treinado, mescla-o com o modelo base e registra o modelo combinado no Unity Catalog via MLflow. Isso torna o modelo disponível para implantação e inferência.

Nota: Este passo requer compute de GPU H100 para carregar o ponto de verificação do modelo. Executar o programa em GPUs menores pode resultar em erros de falta de memória do CUDA.

Python
from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
try:
    from transformers.activations import NewGELUActivation, PytorchGELUTanh, GELUActivation
except ImportError:
    from transformers.activations import NewGELUActivation, GELUTanh as PytorchGELUTanh, GELUActivation

from peft import PeftModel
import mlflow
from mlflow import transformers as mlflow_transformers
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()

components = {
    "model": merged_model,
    "tokenizer": tokenizer,
}

# 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=peft_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,   # 🚨 pass the pipeline, not just the model
        artifact_path="model",
        input_example=input_example,
        # optional: save_pretrained=False for reference-only PEFT logging
        # save_pretrained=False,
    )
# Start MLflow run and log model
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 deployment information
print(f"\n📦 Model Registration Complete!")
print(f"Unity Catalog Path: {full_model_name}")
print(f"Optimization: Liger Kernels + LoRA")

Próximos passos​

Exemplo de caderno​

Ajuste fino do Olmo3 7B com Axolotl em compute serverless múltiplas GPUs.

Abrir notebook em uma nova aba