Ajuste fino distribuído do OpenAI gpt-oss-20b
Este notebook demonstra como ajustar o modelo gpt-oss-20b da OpenAI usando treinamento distribuído em compute GPU serverless . Você aprenderá como:
- Aplique LoRA (Adaptação de Baixa Classificação) para ajustar com eficiência um modelo de 20 bilhões de parâmetros.
- Utilize a quantização MXFP4 para reduzir os requisitos de memória durante o treinamento.
- Aproveite o paralelismo de dados distribuído em 8 GPUs H100.
- Registre o modelo otimizado no Unity Catalog para implantação.
conceitos-chave:
- gpt-oss-20b: Modelo de linguagem de código aberto da OpenAI com 20 bilhões de parâmetros
- LoRa: Ajuste fino com uso eficiente de parâmetros que ensina pequenas camadas adaptadoras enquanto congela o modelo base.
- Quantização MXFP4: Formato de ponto flutuante de 4 bits em microescala que reduz o uso de memória.
- TRL: Biblioteca de Aprendizado por Reforço Transformer para ajuste fino supervisionado
- computeGPU sem servidor: Databricks gerencia compute que escala automaticamente os recursos de GPU.
Conecte-se à computeGPU serverless
Este notebook requer compute GPU serverless . Para conectar:
- Clique no seletor compute do Notebook no canto superior direito e selecione GPU sem servidor.
- No lado direito, clique no botão de ambiente.
- Selecione 8xH100 como acelerador.
- Clique em Aplicar
A função de treinamento provisionará automaticamente 8 GPUs H100 para treinamento distribuído.
Instale a biblioteca necessária.
Instale a biblioteca necessária para treinamento distribuído, incluindo TRL para ajuste fino supervisionado, PEFT para adaptadores LoRa e MLflow para acompanhamento de modelos.
Python
%pip install "trl>=0.20.0" "peft>=0.17.0" "transformers==4.56.1"
%pip install mlflow>=3.6
%pip install hf_transfer==0.1.9
%restart_python
Python
dbutils.library.restartPython()
Configure Unity Catalog e os parâmetros do modelo.
Configure os parâmetros de configuração para o registro Unity Catalog e o treinamento do modelo. Você pode personalizar esses parâmetros usando os widgets acima:
- uc_catalog , uc_schema , uc_model_name : Localização Unity Catalog para registro do modelo
- uc_volume : Nome do volume para armazenar os pontos de verificação do modelo.
- Modelo : Identificador do modelo Hugging Face (default: openai/gpt-oss-20b)
- dataset_path : conjunto de dados a ser usado para ajuste fino (default: HuggingFaceH4/Multilingual-Thinking)
Python
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 dataset
Por default, este Notebook usa 'HuggingFaceH4/Multilingual-Thinking', que foi especificamente selecionado com sequências de raciocínio traduzidas em vários idiomas. Você pode editar o parâmetro "Caminho do conjunto de dados" acima para usar outro dataset.
Defina as utilidades de registro de memória da GPU
Esta função de utilidades ajuda a monitorar o uso da memória da GPU durante o treinamento distribuído. Ele logs a memória alocada e reservada para cada rank da GPU, o que é útil para problemas de memória.
Python
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"
)
Defina a função de treinamento distribuído
A célula seguinte define a função de treinamento usando o decorador @distributed da biblioteca serverless_gpu. Este decorador:
- provisionamento 8 GPUs H100 sob demanda para treinamento distribuído
- Gerencia o paralelismo de dados em várias GPUs automaticamente.
A função inclui:
- carregamento e tokenização de conjuntos de dados
- Inicialização do modelo com quantização MXFP4
- Configuração do adaptador LoRa
- treinamento com checkpointing de gradiente e precisão mista
- Salvando modelos em volumes Unity Catalog
Python
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
Este celular executa a função de treinamento em GPUs 8 H100. O treinamento normalmente leva de 30 a 60 minutos, dependendo do tamanho dataset e da disponibilidade de compute . A função retorna o ID de execução do MLflow para o registro do modelo.
Python
run_id = run_train.distributed()[0]
registrar modelo no Unity Catalog
Agora você pode registrar o modelo otimizado no MLflow e Unity Catalog para implantação.
Importante: Devido ao tamanho do modelo (20B parâmetros), reconecte o Notebook ao acelerador H100 antes de executar as células de registro.
O processo de inscrição irá:
- Carregue o modelo base e merge o com os adaptadores LoRa otimizados.
- Criar um pipeline de geração de texto
- Registre o modelo no MLflow com o registro Unity Catalog
Python
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...")
Verifique o parâmetro de registro
Esta célula verifica o parâmetro register_model . 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.
Python
print("\nRegistering model with MLflow and Unity Catalog...")
from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
from peft import PeftModel
import mlflow
from mlflow import transformers as mlflow_transformers
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 as capacidades de raciocínio multilíngue
O modelo otimizado foi treinado no dataset Multilingual-Thinking, que inclui raciocínio em cadeia em vários idiomas.
A célula a seguir demonstra essa capacidade por meio de:
- Definindo o idioma de raciocínio para alemão.
- Fornecendo uma pergunta em espanhol ("Qual é a capital da Austrália?")
- Observando que o raciocínio interno do modelo é realizado em alemão.
Python
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óximos passos
Agora que você ajustou e testou seu modelo, você pode:
- Implantei o modelo : Modelos de serviço com modelo de instalação
- Saiba mais sobre treinamento distribuído : Treinamento distribuído com múltiplas GPUs e múltiplos nós
- Otimize o uso de GPUs serverless : Melhores práticas para computeem GPU sem servidor.
- Solução de problemas : Solucione problemas em computeGPU serverless
- Saiba mais sobre o gpt-oss da OpenAI : o guia de otimização da OpenAI.