Ajuste fino distribuído do OpenAI gpt-oss-20b

Abra a versão de caderno desta página

Ajuste o modelo gpt-oss-20b do OpenAI com treinamento distribuído no AI Runtime. Você aprenderá a:

  • Aplicar LoRA (Low-Rank Adaptation) para ajustar com eficiência um modelo com 20 bilhões de parâmetros.
  • Usar a quantização MXFP4 para reduzir os requisitos de memória durante o treinamento
  • Aproveitar o paralelismo de dados distribuídos em 8 GPUs H100
  • Registrar o modelo ajustado no Catálogo do Unity para implantação

Principais conceitos:

  • gpt-oss-20b: modelo de linguagem de software livre de parâmetros de 20 bilhões do OpenAI
  • LoRA: Ajuste fino eficiente em termos de parâmetros que treina pequenas camadas adaptadoras enquanto mantém o modelo base congelado
  • Quantização MXFP4: formato de ponto flutuante de microrredimensionamento de 4 bits que reduz o uso de memória
  • TRL: Biblioteca de Aprendizado de Reforço com Transformadores para ajuste supervisionado
  • Runtime de IA: computação gerenciada pela Databricks que automaticamente dimensiona os recursos de unidades de processamento gráfico (GPU)

Conectar-se à computação de GPU sem servidor

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

  1. Clique no seletor de computação do notebook no canto superior direito e selecione GPU sem servidor
  2. No lado direito, clique no botão de ambiente
  3. Selecione 8xH100 como o Acelerador
  4. Escolha o ambiente de IA v5 no painel direito que contém todas as bibliotecas necessárias para executar este exemplo de notebook
  5. Clique em Aplicar

A função de treinamento provisionará automaticamente 8 GPUs H100 para treinamento distribuído.

Configurar parâmetros de catálogo e modelo do Unity

Configure os parâmetros de configuração para o registro do Unity Catalog e o treino de modelo. Você pode personalizar esses parâmetros usando os widgets acima:

  • uc_catalog, uc_schema, uc_model_name: localização do Unity Catalog para registro de modelo
  • uc_volume: Nome do volume para armazenar pontos de verificação de modelo
  • modelo: Identificador de modelo Hugging Face (padrão: openai/gpt-oss-20b)
  • dataset_path: conjunto de dados a ser usado para ajuste fino (padrão: HuggingFaceH4/Multilingual-Thinking)
dbutils.widgets.text("uc_catalog", "main")
dbutils.widgets.text("uc_schema", "default")
dbutils.widgets.text("uc_model_name", "gpt-oss-20b-peft")
dbutils.widgets.text("uc_volume", "checkpoints")
dbutils.widgets.text("model", "openai/gpt-oss-20b")
dbutils.widgets.text("dataset_path", "HuggingFaceH4/Multilingual-Thinking")

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

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

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

Escolha seu conjunto de dados

Por padrão, este notebook usa 'HuggingFaceH4/Multilingual-Thinking', que foi especificamente selecionado com cadeias de pensamento traduzidas em vários idiomas. Você pode editar o parâmetro "Caminho do Conjunto de Dados" acima para usar outro conjunto de dados.

Definir utilitário de monitoramento de memória GPU

Essa função de utilitário ajuda a monitorar o uso de memória gpu durante o treinamento distribuído. Ele registra memória alocada e reservada para cada classificação de GPU, o que é útil para depurar problemas de memória.

import os
import torch
import torch.distributed as dist

def log_gpu_memory(tag=""):
    if not torch.cuda.is_available():
        return

    # rank info (if distributed is initialized)
    if dist.is_available() and dist.is_initialized():
        rank = dist.get_rank()
        world_size = dist.get_world_size()
    else:
        rank = 0
        world_size = 1

    device = torch.cuda.current_device()  # current GPU for this process
    torch.cuda.synchronize(device)

    allocated = torch.cuda.memory_allocated(device) / 1024**2
    reserved  = torch.cuda.memory_reserved(device) / 1024**2

    print(
        f"[{tag}] rank={rank}/{world_size-1}, "
        f"device={device}, "
        f"allocated={allocated:.1f} MB, reserved={reserved:.1f} MB"
    )

Definir a função de treinamento distribuído

A célula a seguir define a função de treinamento usando o decorador @distributed da biblioteca Serverless_GPU. Este decorador:

  • Provisiona 8 GPUs H100 sob demanda para treinamento distribuído
  • Manipula o paralelismo de dados em várias GPUs automaticamente

A função inclui:

  • Carregamento e tokenização do conjunto de dados
  • Inicialização de modelo com quantização MXFP4
  • Configuração do adaptador LoRA
  • Treinamento com checkpoint de gradiente e precisão mista
  • Salvar modelos em volumes do Catálogo Unity
from serverless_gpu import distributed

@distributed(gpus=8, gpu_type="h100")
def run_train():
    import logging
    import os
    import torch

    rank = int(os.environ.get("RANK", "0"))
    local_rank = int(os.environ.get("LOCAL_RANK", "0"))
    torch.cuda.set_device(local_rank)
    world_size = int(os.environ.get("WORLD_SIZE", str(torch.cuda.device_count())))

    os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")

    is_main = rank == 0
    if is_main:
        logging.info("DDP environment")
        logging.info(f"\tWORLD_SIZE={world_size}  RANK={rank}  LOCAL_RANK={local_rank}")
        logging.info(f"\tCUDA device count (this node): {torch.cuda.device_count()}")

    from datasets import load_dataset
    dataset = load_dataset(DATASET_PATH, split="train")

    from transformers import AutoTokenizer
    tokenizer = AutoTokenizer.from_pretrained(HF_MODEL_NAME)

    from transformers import AutoModelForCausalLM, Mxfp4Config

    quantization_config = Mxfp4Config(dequantize=True)
    model_kwargs = dict(
        attn_implementation="eager", # Use eager attention implementation for better performance
        dtype=torch.bfloat16,
        quantization_config=quantization_config,
        use_cache=False, # Since using gradient checkpointing
    )

    model = AutoModelForCausalLM.from_pretrained(HF_MODEL_NAME, **model_kwargs)

    from peft import LoraConfig, get_peft_model

    peft_config = LoraConfig(
        r=8,
        lora_alpha=16,
        target_modules="all-linear",
        lora_dropout=0.05,
        bias="none",
        task_type="CAUSAL_LM",
    )
    peft_model = get_peft_model(model, peft_config)
    if is_main:
        peft_model.print_trainable_parameters()

    from trl import SFTConfig

    training_args = SFTConfig(
        learning_rate=2e-4,
        num_train_epochs=1,
        logging_steps=1,
        per_device_train_batch_size=1,
        gradient_accumulation_steps=2,
        gradient_checkpointing=True,
        gradient_checkpointing_kwargs={"use_reentrant": False},
        max_length=2048,
        warmup_ratio=0.03,
        lr_scheduler_type="cosine_with_min_lr",
        lr_scheduler_kwargs={"min_lr_rate": 0.1},
        output_dir=OUTPUT_DIR,
        report_to="mlflow",  # No reporting to avoid Gradio issues
        push_to_hub=False,  # Disable push to hub to avoid authentication issues
        logging_dir=None,  # Disable tensorboard logging
        disable_tqdm=False,  # Keep progress bars for monitoring
        ddp_find_unused_parameters=False,
    )

    from trl import SFTTrainer

    trainer = SFTTrainer(
        model=peft_model,
        args=training_args,
        train_dataset=dataset,
        processing_class=tokenizer,
    )
    #torch.cuda.empty_cache()
    #log_gpu_memory("before model training")
    result = trainer.train()
    #log_gpu_memory("after model loading")

    if is_main:
        logging.info("Training complete!")
        logging.info(f"Final training loss: {result.training_loss:.4f}")
        logging.info(f"Train runtime (s): {result.metrics.get('train_runtime', 'N/A')}")
        logging.info(f"Samples/sec: {result.metrics.get('train_samples_per_second', 'N/A')}")
        logging.info("\nSaving trained model...")
        trainer.save_model(OUTPUT_DIR)
        logging.info("✓ LoRA adapters saved - use with base model for inference")
        tokenizer.save_pretrained(OUTPUT_DIR)
        logging.info("✓ Tokenizer saved with model")
        logging.info(f"\n🎉 All artifacts saved to: {OUTPUT_DIR}")

    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

Executar o treinamento distribuído

Essa célula executa a função de treinamento em 8 GPUs H100. O treinamento normalmente leva de 30 a 60 minutos, dependendo do tamanho do conjunto de dados e da disponibilidade de computação. A função retorna a ID de execução do MLflow para registro de modelo.

run_id = run_train.distributed()[0]

Registrar modelo no Catálogo do Unity

Agora você pode registrar o modelo ajustado com o MLflow e o Catálogo do Unity para implantação.

Importante: Considerando o tamanho do modelo (parâmetros de 20B), reconecte o notebook ao acelerador H100 antes de executar as células de registro.

O processo de registro será:

  1. Carregue o modelo base e mescle-o com os adaptadores LoRA finamente ajustados
  2. Criar um pipeline de geração de texto
  3. Registre o modelo no MLflow com registro no Unity Catalog
dbutils.widgets.dropdown("register_model", "False", ["True", "False"])
register_model = dbutils.widgets.get("register_model")
if register_model == "False":
  dbutils.notebook.exit("Skipping model registration...")

Verificar o parâmetro de registro

Essa célula verifica o register_model parâmetro. Se definido como False, o notebook ignorará o registro do modelo. Você pode alterar esse parâmetro usando o widget na parte superior do notebook.

print("\nRegistering model with MLflow and Unity Catalog...")

from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline

from peft import PeftModel
import mlflow
import torch

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

Testar recursos de raciocínio multilíngue

O modelo finamente ajustado foi treinado no conjunto de dados Multilingual-Thinking, que inclui raciocínio encadeado em vários idiomas.

A célula a seguir demonstra essa funcionalidade:

  • Definindo o idioma de raciocínio como alemão
  • Fornecendo um prompt em espanhol ("Qual é a capital da Austrália?")
  • Observando que o raciocínio interno do modelo é executado em alemão
REASONING_LANGUAGE = "German"
SYSTEM_PROMPT = f"reasoning language: {REASONING_LANGUAGE}"
USER_PROMPT = "¿Cuál es el capital de Australia?"  # Spanish for "What is the capital of Australia?"

messages = [
    {"role": "system", "content": SYSTEM_PROMPT},
    {"role": "user", "content": USER_PROMPT},
]

input_ids = tokenizer.apply_chat_template(
    messages,
    add_generation_prompt=True,
    return_tensors="pt",
).to(merged_model.device)

gen_kwargs = {"max_new_tokens": 512, "do_sample": True, "temperature": 0.6, "top_p": None, "top_k": None}

output_ids = merged_model.generate(input_ids, **gen_kwargs)
response = tokenizer.batch_decode(output_ids)[0]
print(response)

Próximas Etapas 

Agora que você ajustou e testou seu modelo, você pode:

Notebook de exemplo

Ajuste fino distribuído do OpenAI gpt-oss-20b

Obter laptop

  • Last updated on