Dopasowanie Llama 3.1 8B przy użyciu Mosaic LLM Foundry na bezserwerowym GPU Databricks.

Otwórz wersję notebooka tej strony

Dostrajanie modelu Llama 3.1 8B w środowisku uruchomieniowym sztucznej inteligencji przy użyciu narzędzia Mosaic LLM Foundry, bazy kodu do trenowania, dostrajania, oceniania i wdrażania dużych modeli językowych z obsługą strategii trenowania rozproszonego.

Notatnik używa funkcji:

  • Mozaika LLM Foundry: struktura do trenowania i dostrajania LLM z wbudowaną obsługą FSDP, wydajne ładowanie danych i integracja MLflow
  • FSDP (w pełni fragmentowane dane równoległe): dystrybuuje parametry modelu, gradienty i stany optymalizatora między procesorami GPU
  • Bezserwerowe GPU Databricks: umożliwia rozproszone trenowanie na połączonych bezserwerowych zasobach obliczeniowych GPU
  • Unity Catalog: przechowuje punkty kontrolne i rejestruje trenowane modele
  • MLflow: śledzi eksperymenty i rejestruje metryki trenowania

Note

Ten przykład wymaga standardowego środowiska w wersji 5 lub wyższej.

Nawiązywanie połączenia z bezserwerowym przetwarzaniem GPU

Ten notebook wymaga przetwarzania bezserwerowego GPU. Aby nawiązać połączenie:

  1. Kliknij selektor obliczeniowy notesu w prawym górnym rogu i wybierz pozycję Bezserwerowy procesor GPU
  2. Otwórz panel boczny Środowisko po prawej stronie notesu
  3. Ustaw akcelerator na 8xH100
  4. Wybierz środowisko podstawowe w warstwie Standardowa i ustaw wersję środowiska na 5, która zawiera biblioteki potrzebne do uruchomienia tego przykładu
  5. Wybierz przycisk Zastosuj i kliknij przycisk Potwierdź, aby zastosować to środowisko do notatnika

Instalowanie wymaganych bibliotek

Zainstaluj oprogramowanie Mosaic LLM Foundry i jego zależności na potrzeby trenowania rozproszonego. Wstępnie zbudowany pakiet flash-attn wheel jest instalowany jako pierwszy, aby pip wykorzystał go ponownie zamiast kompilować flash-attention ze źródeł (co jest powolne) podczas rozwiązywania zależności llm-foundry[gpu]:

  • flash-attn: Zoptymalizowana implementacja mechanizmu attention, zainstalowana z prekompilowanego pakietu wheel
  • llm-foundry: Podstawowa platforma do trenowania i dostrajania modeli LLM
  • hf_transfer: Szybsze pobieranie modeli z Hugging Face
  • yamlmagic: Włącza konfigurację YAML w komórkach notatnika
%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

Uruchom ponownie środowisko języka Python

Uruchom ponownie jądro języka Python, aby upewnić się, że nowo zainstalowane pakiety są dostępne.

dbutils.library.restartPython()

Konfigurowanie ścieżek Unity Catalog dla przechowywania modeli

Skonfiguruj lokalizacje Unity Catalog do przechowywania checkpointów modelu i rejestrowania wytrenowanego modelu. Konfiguracja używa parametrów zapytania, które można dostosować bez edytowania kodu.

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

Definiowanie konfiguracji trenowania przy użyciu języka YAML

Załaduj konfigurację dostrajania z formatu YAML. Konfiguracja określa:

  • Architektura modelu i wstępnie wytrenowane parametry (Llama 3.1 8B)
  • Ustawienia FSDP na potrzeby trenowania rozproszonego
  • Hiperparametry trenowania (szybkość uczenia, rozmiar partii, optymalizator)
  • Konfiguracja zestawu danych (MosaicML/dolly_hhrlhf)
  • Rejestrowanie i modelowanie punktów kontrolnych platformy MLflow
  • Wywołania zwrotne na potrzeby monitorowania i optymalizacji
%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}"

Zdefiniuj funkcję trenowania rozproszonego

Ta komórka definiuje funkcję trenowania, która zostanie uruchomiona na 8 układach GPU H100 przy użyciu dekoratora @distributed. Funkcja:

  • Konfiguruje token Hugging Face w celu uzyskania dostępu do modelu
  • Umożliwia szybkie pobieranie modeli za pomocą polecenia hf_transfer
  • Wywołuje funkcję LLM Foundry train() z konfiguracją YAML
  • Zwraca identyfikator przebiegu MLflow do śledzenia eksperymentu

Dekorator @distributed uruchamia funkcję na połączonym bezserwerowym procesorze GPU obliczeniowym i obsługuje rozproszoną aranżację trenowania.

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

Uruchom rozproszoną pracę treningową

Wykonaj funkcję szkoleniową na 8 procesorach GPU H100. Funkcja zwraca identyfikator przebiegu platformy MLflow, który może służyć do śledzenia metryk, wyświetlania dzienników i uzyskiwania dostępu do wytrenowanego modelu w interfejsie użytkownika platformy MLflow.

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

Następne kroki

Przykładowy notatnik

Dopasowanie Llama 3.1 8B przy użyciu Mosaic LLM Foundry na bezserwerowym GPU Databricks.

Pobierz laptopa

  • Last updated on