Stel Llama 3.1 8B af met behulp van Mosaic LLM Foundry op Databricks Serverless GPU.

Notitieblokversie van deze pagina openen

Verfijn een Llama 3.1 8B-model op AI Runtime met behulp van Mosaic LLM Foundry, een codebasis voor training, verfijning, evaluatie en implementatie van grote taalmodellen met ondersteuning voor gedistribueerde trainingsstrategieën.

In het notebook wordt het volgende gebruikt:

  • Mozaïek LLM Foundry: Een framework voor het trainen en verfijnen van LLM's met ingebouwde ondersteuning voor FSDP, efficiënt laden van gegevens en MLflow-integratie
  • FSDP (Volledig Sharded Data Parallel): hiermee worden modelparameters, gradiënten en optimizerstatussen verdeeld over GPU's
  • Databricks Serverless GPU: voert gedistribueerde training uit op verbonden serverloze GPU-rekenkracht
  • Unity Catalog: Modelcontrolepunten opslaan en getrainde modellen registreren
  • MLflow: houdt experimenten bij en registreert metrische gegevens voor training

Verbinding maken met serverloze GPU-rekenkracht

Voor dit notebook is serverloze GPU-rekenkracht vereist. Verbinding maken:

  1. Klik op de rekenkiezer van het notebook in de rechterbovenhoek en selecteer serverloze GPU
  2. Open het deelvenster Omgeving aan de rechterkant van het notitieblok
  3. Accelerator instellen op 8xH100
  4. Selecteer de standaardbasisomgeving en stel omgevingsversiein op 5, die de bibliotheken bevat die nodig zijn om dit voorbeeld uit te voeren
  5. Selecteer Toepassen en klik op Bevestigen om deze omgeving toe te passen op uw notitieblok

Vereiste bibliotheken installeren

Installeer Mosaic LLM Foundry en de bijbehorende afhankelijkheden voor gedistribueerde training. Het vooraf gebouwde flash-attn wiel wordt eerst geïnstalleerd, zodat pip het opnieuw wordt gebruikt in plaats van flash-aandacht te compileren vanuit de bron (wat traag is) wanneer het wordt omgezet llm-foundry[gpu]:

  • flash-attn: Geoptimaliseerde aandachtimplementatie, geïnstalleerd vanaf een vooraf gebouwd wiel
  • llm-foundry: Kernframework voor LLM-training en afstemming
  • hf_transfer: Snellere downloads van modellen van Hugging Face
  • yamlmagic: Hiermee schakelt u YAML-configuratie in notebookcellen in
%pip install --no-deps "https://github.com/Dao-AILab/flash-attention/releases/download/v2.7.4.post1/flash_attn-2.7.4.post1+cu12torch2.6cxx11abiFALSE-cp312-cp312-linux_x86_64.whl"
%pip install llm-foundry[gpu]==0.20.0
%pip install hf_transfer
%pip install git+https://github.com/josejg/yamlmagic.git

De Python-omgeving opnieuw starten

Start de Python-kernel opnieuw om ervoor te zorgen dat de zojuist geïnstalleerde pakketten beschikbaar zijn.

dbutils.library.restartPython()

Unity Catalog-paden configureren voor modelopslag

Stel Unity Catalog-locaties in voor het opslaan van modelcontrolepunten en het registreren van het getrainde model. De configuratie maakt gebruik van queryparameters die kunnen worden aangepast zonder de code te bewerken.

dbutils.widgets.text("uc_catalog", "main")
dbutils.widgets.text("uc_schema", "default")
dbutils.widgets.text("uc_model_name", "llama3_1-8b")
dbutils.widgets.text("uc_volume", "checkpoints")

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

MLFLOW_EXPERIMENT_NAME = '/Workspace/Shared/llm-foundry-sgc' # TODO: update this name

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"EXPERIMENT_NAME: {MLFLOW_EXPERIMENT_NAME}")

# Model selection - Choose based on your compute constraints
OUTPUT_DIR = f"/Volumes/{UC_CATALOG}/{UC_SCHEMA}/{UC_VOLUME}/{UC_MODEL_NAME}" # Save checkpoint to UC Volume

print(f"OUTPUT_DIR: {OUTPUT_DIR}")

Trainingsconfiguratie definiëren met YAML

Laad de fijn-afstemmingsconfiguratie vanuit het YAML-formaat. De configuratie geeft het volgende op:

  • Modelarchitectuur en vooraf getrainde gewichten (Llama 3.1 8B)
  • FSDP-instellingen voor gedistribueerde training
  • Hyperparameters trainen (leersnelheid, batchgrootte, optimizer)
  • Gegevenssetconfiguratie (mozaïekml/dolly_hhrlhf)
  • MLflow-logboekregistratie en modelcontrolepunten
  • Callbacks voor bewaking en optimalisatie
%load_ext yamlmagic

%%yaml config
seed: 17
model:
  name: hf_causal_lm
  pretrained: true
  init_device: mixed
  use_auth_token: true
  use_flash_attention_2: true
  pretrained_model_name_or_path: meta-llama/Llama-3.1-8B
loggers:
  mlflow:
    resume: true
    tracking_uri: databricks
    rename_metrics:
      time/token: time/num_tokens
      lr-DecoupledLionW/group0: learning_rate
    log_system_metrics: true
    experiment_name: "mlflow_experiment_name"
    run_name: llama3_8b-finetune
    model_registry_uri: databricks-uc
    model_registry_prefix: main.linyuan
callbacks:
  lr_monitor: {}
  run_timeout:
    timeout: 7200
  scheduled_gc:
    batch_interval: 1000
  speed_monitor:
    window_size: 10
  memory_monitor: {}
  runtime_estimator: {}
  hf_checkpointer:
    save_folder: "dbfs:/Volumes/main/sgc/checkpoints/llama3_1-8b-hf"
    save_interval: "1ep"
    precision: "bfloat16"
    overwrite: true

    mlflow_registered_model_name: "main.sgc.llama3_1_8b_full_ft"
    mlflow_logging_config:
      task: "llm/v1/completions"
      metadata:
        pretrained_model_name: "meta-llama/Llama-3.1-8B-Instruct"
optimizer:
  lr: 5.0e-07
  name: decoupled_lionw
  betas:
  - 0.9
  - 0.95
  weight_decay: 0
precision: amp_bf16
scheduler:
  name: linear_decay_with_warmup
  alpha_f: 0
  t_warmup: 10ba
tokenizer:
  name: meta-llama/Llama-3.1-8B
  kwargs:
    model_max_length: 1024
algorithms:
  gradient_clipping:
    clipping_type: norm
    clipping_threshold: 1
autoresume: false
log_config: false
fsdp_config:
  verbose: false
  mixed_precision: PURE
  state_dict_type: sharded
  limit_all_gathers: true
  sharding_strategy: FULL_SHARD
  activation_cpu_offload: false
  activation_checkpointing: true
  activation_checkpointing_reentrant: false
max_seq_len: 1024
save_folder: "output_folder"
dist_timeout: 600
max_duration: 20ba
progress_bar: false
train_loader:
  name: finetuning
  dataset:
    split: test
    hf_name: mosaicml/dolly_hhrlhf
    shuffle: true
    safe_load: true
    max_seq_len: 1024
    packing_ratio: auto
    target_prompts: none
    target_responses: all
    allow_pad_trimming: false
    decoder_only_format: true
  timeout: 0
  drop_last: false
  pin_memory: true
  num_workers: 8
  prefetch_factor: 2
  persistent_workers: true
eval_interval: 1
save_interval: 1h
log_to_console: true
save_overwrite: true
python_log_level: debug
save_weights_only: false
console_log_interval: 10ba
device_eval_batch_size: 1
global_train_batch_size: 32
device_train_microbatch_size: 1
save_num_checkpoints_to_keep: 1

config["loggers"]["mlflow"]["experiment_name"] = MLFLOW_EXPERIMENT_NAME
config["save_folder"] = OUTPUT_DIR
config["callbacks"]["hf_checkpointer"]["save_folder"] = OUTPUT_DIR
config["callbacks"]["hf_checkpointer"]["mlflow_registered_model_name"] = f"{UC_CATALOG}.{UC_SCHEMA}.{UC_MODEL_NAME}"

De gedistribueerde trainingsfunctie definiëren

Deze cel definieert de trainingsfunctie die wordt uitgevoerd op 8 H100 GPU's met behulp van de @distributed decorator. De functie:

  • Hiermee configureert u het Hugging Face-token voor modeltoegang
  • Maakt snelle modeldownloads mogelijk met hf_transfer
  • Roept de functie LLM Foundry train() aan met de YAML-configuratie
  • Retourneert de MLflow-uitvoerings-id voor het bijhouden van het experiment

De @distributed-decorator voert de functie uit op de verbonden serverloze GPU-rekenomgeving en verzorgt de orkestratie van gedistribueerde training.

from serverless_gpu import distributed
from llmfoundry.command_utils.train import train
from omegaconf import DictConfig
import mlflow
from huggingface_hub import constants

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

@distributed(gpus=8, gpu_type='H100')
def run_llm_foundry():
    import os
    import logging
    os.environ["HUGGING_FACE_HUB_TOKEN"] = HF_TOKEN
    constants.HF_HUB_ENABLE_HF_TRANSFER = True
    train(DictConfig(config))

    logging.info("\n✓ Training completed successfully!")

    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

De gedistribueerde trainingstaak uitvoeren

Voer de trainingsfunctie uit op 8 H100 GPU's. De functie retourneert de MLflow-uitvoerings-id, die kan worden gebruikt om metrische gegevens bij te houden, logboeken weer te geven en toegang te krijgen tot het getrainde model in de MLflow-gebruikersinterface.

mlflow_run_id = run_llm_foundry.distributed()[0]
print(mlflow_run_id)

Volgende stappen 

Voorbeeld van notebook

Stel Llama 3.1 8B af met behulp van Mosaic LLM Foundry op Databricks Serverless GPU.

Notebook krijgen

  • Last updated on