mirror of
https://github.com/meta-llama/llama-stack.git
synced 2025-07-26 22:19:49 +00:00
279 lines
9.8 KiB
Python
279 lines
9.8 KiB
Python
# Copyright (c) Meta Platforms, Inc. and affiliates.
|
|
# All rights reserved.
|
|
#
|
|
# This source code is licensed under the terms described in the LICENSE file in
|
|
# the root directory of this source tree.
|
|
|
|
import logging
|
|
import multiprocessing
|
|
import os
|
|
import signal
|
|
import sys
|
|
from datetime import UTC, datetime
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
import psutil
|
|
import torch
|
|
from datasets import Dataset
|
|
from transformers import AutoConfig, AutoModelForCausalLM
|
|
|
|
from llama_stack.apis.datasetio import DatasetIO
|
|
from llama_stack.apis.post_training import Checkpoint, TrainingConfig
|
|
|
|
from .config import HuggingFacePostTrainingConfig
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
def setup_environment():
|
|
"""Setup common environment variables for training."""
|
|
os.environ["TOKENIZERS_PARALLELISM"] = "false"
|
|
os.environ["MKL_THREADING_LAYER"] = "GNU"
|
|
os.environ["MKL_SERVICE_FORCE_INTEL"] = "0"
|
|
os.environ["MKL_NUM_THREADS"] = "1"
|
|
|
|
|
|
def get_gb(to_convert: int) -> str:
|
|
"""Converts memory stats to GB and formats to 2 decimal places.
|
|
Args:
|
|
to_convert: Memory value in bytes
|
|
Returns:
|
|
str: Memory value in GB formatted to 2 decimal places
|
|
"""
|
|
return f"{(to_convert / (1024**3)):.2f}"
|
|
|
|
|
|
def get_memory_stats(device: torch.device) -> dict[str, Any]:
|
|
"""Get memory statistics for the given device."""
|
|
stats = {
|
|
"system_memory": {
|
|
"total": get_gb(psutil.virtual_memory().total),
|
|
"available": get_gb(psutil.virtual_memory().available),
|
|
"used": get_gb(psutil.virtual_memory().used),
|
|
"percent": psutil.virtual_memory().percent,
|
|
}
|
|
}
|
|
|
|
if device.type == "cuda":
|
|
stats["device_memory"] = {
|
|
"allocated": get_gb(torch.cuda.memory_allocated(device)),
|
|
"reserved": get_gb(torch.cuda.memory_reserved(device)),
|
|
"max_allocated": get_gb(torch.cuda.max_memory_allocated(device)),
|
|
}
|
|
elif device.type == "mps":
|
|
# MPS doesn't provide direct memory stats, but we can track system memory
|
|
stats["device_memory"] = {
|
|
"note": "MPS memory stats not directly available",
|
|
"system_memory_used": get_gb(psutil.virtual_memory().used),
|
|
}
|
|
elif device.type == "cpu":
|
|
# For CPU, we track process memory usage
|
|
process = psutil.Process()
|
|
stats["device_memory"] = {
|
|
"process_rss": get_gb(process.memory_info().rss),
|
|
"process_vms": get_gb(process.memory_info().vms),
|
|
"process_percent": process.memory_percent(),
|
|
}
|
|
|
|
return stats
|
|
|
|
|
|
def setup_torch_device(device_str: str) -> torch.device:
|
|
"""Initialize and validate a PyTorch device.
|
|
This function handles device initialization and validation for different device types:
|
|
- CUDA: Validates CUDA availability and handles device selection
|
|
- MPS: Validates MPS availability for Apple Silicon
|
|
- CPU: Basic validation
|
|
- HPU: Raises error as it's not supported
|
|
Args:
|
|
device_str: String specifying the device ('cuda', 'cpu', 'mps')
|
|
Returns:
|
|
torch.device: The initialized and validated device
|
|
Raises:
|
|
RuntimeError: If device initialization fails or device is not supported
|
|
"""
|
|
try:
|
|
device = torch.device(device_str)
|
|
except RuntimeError as e:
|
|
raise RuntimeError(f"Error getting Torch Device {str(e)}") from e
|
|
|
|
# Validate device capabilities
|
|
if device.type == "cuda":
|
|
if not torch.cuda.is_available():
|
|
raise RuntimeError(
|
|
f"{device.type}: Torch has no CUDA/ROCm support or could not detect a compatible device."
|
|
)
|
|
if device.index is None:
|
|
device = torch.device(device.type, torch.cuda.current_device())
|
|
elif device.type == "mps":
|
|
if not torch.backends.mps.is_available():
|
|
raise RuntimeError(f"{device.type}: Torch has no MPS support or could not detect a compatible device.")
|
|
elif device.type == "hpu":
|
|
raise RuntimeError(f"{device.type}: training does not support Intel Gaudi.")
|
|
|
|
return device
|
|
|
|
|
|
async def setup_data(datasetio_api: DatasetIO, dataset_id: str) -> list[dict[str, Any]]:
|
|
"""Load dataset from llama stack dataset provider"""
|
|
try:
|
|
all_rows = await datasetio_api.iterrows(
|
|
dataset_id=dataset_id,
|
|
limit=-1,
|
|
)
|
|
if not isinstance(all_rows.data, list):
|
|
raise RuntimeError("Expected dataset data to be a list")
|
|
return all_rows.data
|
|
except Exception as e:
|
|
raise RuntimeError(f"Failed to load dataset: {str(e)}") from e
|
|
|
|
|
|
def load_model(
|
|
model: str,
|
|
device: torch.device,
|
|
provider_config: HuggingFacePostTrainingConfig,
|
|
) -> AutoModelForCausalLM:
|
|
"""Load and initialize the model for training.
|
|
Args:
|
|
model: The model identifier to load
|
|
device: The device to load the model onto
|
|
provider_config: Provider-specific configuration
|
|
Returns:
|
|
The loaded and initialized model
|
|
Raises:
|
|
RuntimeError: If model loading fails
|
|
"""
|
|
logger.info("Loading the base model")
|
|
try:
|
|
model_config = AutoConfig.from_pretrained(model, **provider_config.model_specific_config)
|
|
model_obj = AutoModelForCausalLM.from_pretrained(
|
|
model,
|
|
torch_dtype="auto" if device.type != "cpu" else "float32",
|
|
quantization_config=None,
|
|
config=model_config,
|
|
**provider_config.model_specific_config,
|
|
)
|
|
# Always move model to specified device
|
|
model_obj = model_obj.to(device)
|
|
logger.info(f"Model loaded and moved to device: {model_obj.device}")
|
|
return model_obj
|
|
except Exception as e:
|
|
raise RuntimeError(f"Failed to load model: {str(e)}") from e
|
|
|
|
|
|
def split_dataset(ds: Dataset) -> tuple[Dataset, Dataset]:
|
|
"""Split dataset into train and validation sets.
|
|
Args:
|
|
ds: Dataset to split
|
|
Returns:
|
|
tuple: (train_dataset, eval_dataset)
|
|
"""
|
|
logger.info("Splitting dataset into train and validation sets")
|
|
train_val_split = ds.train_test_split(test_size=0.1, seed=42)
|
|
train_dataset = train_val_split["train"]
|
|
eval_dataset = train_val_split["test"]
|
|
logger.info(f"Split dataset into {len(train_dataset)} training and {len(eval_dataset)} validation examples")
|
|
return train_dataset, eval_dataset
|
|
|
|
|
|
def setup_signal_handlers():
|
|
"""Setup signal handlers for graceful shutdown."""
|
|
|
|
def signal_handler(signum, frame):
|
|
logger.info(f"Received signal {signum}, initiating graceful shutdown")
|
|
sys.exit(0)
|
|
|
|
signal.signal(signal.SIGTERM, signal_handler)
|
|
signal.signal(signal.SIGINT, signal_handler)
|
|
|
|
|
|
def calculate_training_steps(steps_per_epoch: int, config: TrainingConfig) -> dict[str, int]:
|
|
"""Calculate training steps and logging configuration.
|
|
Args:
|
|
steps_per_epoch: Number of training steps per epoch
|
|
config: Training configuration
|
|
Returns:
|
|
dict: Dictionary with calculated step values
|
|
"""
|
|
total_steps = steps_per_epoch * config.n_epochs
|
|
max_steps = min(config.max_steps_per_epoch, total_steps)
|
|
logging_steps = max(1, steps_per_epoch // 50) # Log 50 times per epoch
|
|
|
|
logger.info("Training configuration:")
|
|
logger.info(f"- Steps per epoch: {steps_per_epoch}")
|
|
logger.info(f"- Total steps: {total_steps}")
|
|
logger.info(f"- Max steps: {max_steps}")
|
|
logger.info(f"- Logging steps: {logging_steps}")
|
|
|
|
return {"total_steps": total_steps, "max_steps": max_steps, "logging_steps": logging_steps}
|
|
|
|
|
|
def get_save_strategy(output_dir_path: Path | None) -> tuple[str, str]:
|
|
"""Get save and evaluation strategy based on output directory.
|
|
Args:
|
|
output_dir_path: Optional path to save the model
|
|
Returns:
|
|
tuple: (save_strategy, eval_strategy)
|
|
"""
|
|
if output_dir_path:
|
|
logger.info(f"Will save checkpoints to {output_dir_path}")
|
|
return "epoch", "epoch"
|
|
return "no", "no"
|
|
|
|
|
|
def create_checkpoints(
|
|
output_dir_path: Path, job_uuid: str, model: str, config: TrainingConfig, final_model_name: str
|
|
) -> list[Checkpoint]:
|
|
"""Create checkpoint objects from training output.
|
|
Args:
|
|
output_dir_path: Path to the training output directory
|
|
job_uuid: Unique identifier for the training job
|
|
model: Model identifier
|
|
config: Training configuration
|
|
final_model_name: Name of the final model directory ("merged_model" for SFT, "dpo_model" for DPO)
|
|
Returns:
|
|
List of Checkpoint objects
|
|
"""
|
|
checkpoints = []
|
|
|
|
# Add checkpoint directories
|
|
checkpoint_dirs = sorted(
|
|
[d for d in output_dir_path.glob("checkpoint-*") if d.is_dir()],
|
|
key=lambda x: int(x.name.split("-")[1]),
|
|
)
|
|
|
|
for epoch_number, checkpoint_dir in enumerate(checkpoint_dirs, start=1):
|
|
created_time = datetime.fromtimestamp(os.path.getctime(checkpoint_dir), tz=UTC)
|
|
checkpoint = Checkpoint(
|
|
identifier=checkpoint_dir.name,
|
|
created_at=created_time,
|
|
epoch=epoch_number,
|
|
post_training_job_id=job_uuid,
|
|
path=str(checkpoint_dir),
|
|
)
|
|
checkpoints.append(checkpoint)
|
|
|
|
# Add final model
|
|
final_model_path = output_dir_path / final_model_name
|
|
if final_model_path.exists():
|
|
training_type = "sft" if final_model_name == "merged_model" else "dpo"
|
|
checkpoint = Checkpoint(
|
|
identifier=f"{model}-{training_type}-{config.n_epochs}",
|
|
created_at=datetime.now(UTC),
|
|
epoch=config.n_epochs,
|
|
post_training_job_id=job_uuid,
|
|
path=str(final_model_path),
|
|
)
|
|
checkpoints.append(checkpoint)
|
|
|
|
return checkpoints
|
|
|
|
|
|
def setup_multiprocessing_for_device(device: torch.device):
|
|
"""Setup multiprocessing start method based on device type.
|
|
Args:
|
|
device: The device being used for training
|
|
"""
|
|
if device.type in ["cuda", "mps"]:
|
|
multiprocessing.set_start_method("spawn", force=True)
|