Skip to content

API_REFERENCE.md

Latest commit

 

History

History
440 lines (327 loc) · 13.9 KB

API_REFERENCE.md

File metadata and controls

440 lines (327 loc) · 13.9 KB
Raw

LM-Trainer API Reference

This document provides a comprehensive reference for all public APIs in the LM-Trainer library.

Package Structure

lm_trainer/
├── cli/              # Command-line interface tools
├── core/             # Core utilities, exceptions, and constants
├── data/             # Data loading and preprocessing
├── extensions/       # Third-party integrations (PEFT, Accelerate, etc.)
├── inference/        # Text generation and model serving
├── models/           # Model architectures and components
├── tokenizers/       # Tokenization utilities
├── training/         # Training loop, optimizers, and evaluation
└── utilities/        # Logging, monitoring, and system utilities

Core Module

Exceptions

LMTrainerError

Base exception for all LM-Trainer errors.

ConfigurationError

Raised when there is an error in configuration.

TrainingError

Raised when there is an error during training.

ValidationError

Raised when there is a validation error.

SerializationError

Raised when there is an error in model serialization.

Constants

DEFAULT_DEVICE

Default device setting ("auto").

DEFAULT_DTYPE

Default data type (torch.float32).

Training Module

Configuration

TrainingConfig

Enhanced configuration for the training process with TRL-style parameters.

Attributes:

  • batch_size: Training batch size (int, default: 32)
  • learning_rate: Learning rate (float, default: 1e-4)
  • weight_decay: Weight decay coefficient (float, default: 0.01)
  • num_epochs: Number of training epochs (int, default: 10)
  • max_steps: Maximum number of training steps (Optional[int], default: None)
  • lr_scheduler: Learning rate scheduler type (str, default: "cosine")
  • warmup_steps: Number of warmup steps (int, default: 1000)
  • warmup_ratio: Ratio of warmup steps (float, default: 0.1)
  • min_lr_ratio: Minimum learning rate ratio (float, default: 0.1)
  • optimizer: Optimizer type (str, default: "adamw")
  • beta1: Beta1 parameter for Adam optimizers (float, default: 0.9)
  • beta2: Beta2 parameter for Adam optimizers (float, default: 0.999)
  • eps: Epsilon parameter for optimizers (float, default: 1e-8)
  • max_grad_norm: Maximum gradient norm for clipping (float, default: 1.0)
  • gradient_accumulation_steps: Gradient accumulation steps (int, default: 1)
  • use_amp: Use automatic mixed precision (bool, default: True)
  • amp_dtype: AMP data type (str, default: "float16")
  • save_steps: Save checkpoint every N steps (int, default: 1000)
  • save_total_limit: Maximum number of checkpoints to keep (int, default: 3)
  • checkpoint_dir: Checkpoint directory (str, default: "./checkpoints")
  • resume_from_checkpoint: Path to resume from (Optional[str], default: None)
  • eval_steps: Evaluation steps (int, default: 500)
  • eval_strategy: Evaluation strategy (str, default: "steps")
  • eval_accumulation_steps: Evaluation accumulation steps (Optional[int], default: None)
  • logging_steps: Logging steps (int, default: 100)
  • log_level: Log level (str, default: "info")
  • report_to: Reporting backends (List[str], default: ["tensorboard"])
  • dataloader_num_workers: Number of data loader workers (int, default: 4)
  • dataloader_pin_memory: Pin memory in data loader (bool, default: True)
  • dataloader_drop_last: Drop last batch in data loader (bool, default: True)
  • local_rank: Local rank for distributed training (int, default: -1)
  • world_size: World size for distributed training (int, default: 1)
  • distributed_backend: Distributed backend (str, default: "nccl")
  • use_deepspeed: Use DeepSpeed (bool, default: False)
  • deepspeed_config: DeepSpeed configuration path (Optional[str], default: None)
  • use_accelerate: Use Accelerate (bool, default: False)
  • accelerate_mixed_precision: Accelerate mixed precision (str, default: "no")
  • use_peft: Use Parameter-Efficient Fine-Tuning (bool, default: False)
  • peft_type: PEFT type (str, default: "lora")
  • peft_r: LoRA rank (int, default: 8)
  • peft_alpha: LoRA alpha (int, default: 16)
  • peft_dropout: LoRA dropout (float, default: 0.05)
  • peft_target_modules: Target modules for PEFT (Optional[List[str]], default: None)
  • peft_bias: LoRA bias setting (str, default: "none")
  • peft_task_type: PEFT task type (str, default: "CAUSAL_LM")
  • device: Device to use (str, default: "auto")
  • force_cpu: Force CPU usage (bool, default: False)
  • seed: Random seed (int, default: 42)
  • early_stopping_patience: Early stopping patience (Optional[int], default: None)
  • early_stopping_threshold: Early stopping threshold (float, default: 0.0)
  • compile_model: Compile model (bool, default: False)
  • compile_mode: Compilation mode (str, default: "default")
  • auto_tune_batch_size: Auto-tune batch size (bool, default: False)
  • use_gradient_compression: Use gradient compression (bool, default: False)
  • gradient_as_bucket_view: Gradient as bucket view (bool, default: True)
  • static_graph: Static graph optimization (bool, default: False)

Methods:

  • get_effective_device(): Get the effective device
  • should_use_amp(): Determine if AMP should be used
  • get_effective_distributed_backend(): Get effective distributed backend
  • save(path): Save configuration to file
  • load(path): Load configuration from file
  • to_dict(): Convert to dictionary
  • from_dict(config_dict): Create from dictionary

ModelConfig

Configuration for Transformer model architecture.

Attributes:

  • vocab_size: Vocabulary size (int, default: 50000)
  • d_model: Model dimension (int, default: 512)
  • n_heads: Number of attention heads (int, default: 8)
  • n_layers: Number of layers (int, default: 6)
  • d_ff: Feed-forward dimension (int, default: 2048)
  • max_seq_len: Maximum sequence length (int, default: 1024)
  • dropout: Dropout rate (float, default: 0.1)
  • attention_dropout: Attention dropout rate (float, default: 0.1)
  • init_std: Initialization standard deviation (float, default: 0.02)
  • use_learned_pos_emb: Use learned positional embeddings (bool, default: False)
  • layer_norm_eps: Layer normalization epsilon (float, default: 1e-5)
  • pre_norm: Use pre-normalization (bool, default: True)
  • activation: Activation function (str, default: "gelu")
  • use_bias: Use bias terms (bool, default: True)
  • gradient_checkpointing: Use gradient checkpointing (bool, default: False)

Properties:

  • d_head: Dimension of each attention head

Methods:

  • save(path): Save configuration to file
  • load(path): Load configuration from file
  • to_dict(): Convert to dictionary
  • from_dict(config_dict): Create from dictionary

Core

Trainer

Base trainer class for training language models.

Methods:

  • __init__(model, train_dataloader, eval_dataloader, optimizer, scheduler, config, model_config, callbacks): Initialize trainer
  • train(): Run the training loop
  • evaluate(): Evaluate the model
  • save_checkpoint(metrics): Save a checkpoint
  • load_checkpoint(checkpoint_path): Load a checkpoint

Callback

Abstract base class for training callbacks.

Methods:

  • on_train_start(trainer): Called at the beginning of training
  • on_train_end(trainer, results): Called at the end of training
  • on_train_error(trainer, error): Called when training encounters an error
  • on_epoch_start(trainer, epoch): Called at the beginning of each epoch
  • on_epoch_end(trainer, epoch, metrics): Called at the end of each epoch
  • on_epoch_begin(trainer, epoch): Called at the beginning of epoch training phase
  • on_epoch_finish(trainer, epoch): Called at the end of epoch training phase
  • on_batch_start(trainer, batch_idx, batch): Called at the beginning of each training batch
  • on_batch_end(trainer, batch_idx, batch, outputs): Called at the end of each training batch
  • on_step_end(trainer, step): Called at the end of each training step
  • on_eval_start(trainer): Called at the beginning of evaluation
  • on_eval_end(trainer, metrics): Called at the end of evaluation
  • on_eval_batch_start(trainer, batch_idx, batch): Called at the beginning of each evaluation batch
  • on_eval_batch_end(trainer, batch_idx, batch, outputs): Called at the end of each evaluation batch

ProgressBarCallback

Simple progress bar callback for training visualization.

TensorBoardCallback

TensorBoard logging callback for training metrics.

EarlyStoppingCallback

Early stopping callback to prevent overfitting.

LearningRateSchedulerCallback

Learning rate scheduler callback.

ModelCheckpointCallback

Model checkpoint callback for saving best models.

CheckpointManager

Manages model checkpoints during training.

Methods:

  • __init__(checkpoint_dir, max_checkpoints): Initialize checkpoint manager
  • save_checkpoint(data, step, is_best, prefix): Save a checkpoint
  • load_checkpoint(checkpoint_path): Load a checkpoint
  • get_latest_checkpoint(): Get the latest checkpoint
  • get_best_checkpoint(): Get the best checkpoint
  • save_training_state(...): Save complete training state
  • load_training_state(checkpoint_path): Load complete training state

Optimization

get_optimizer

Create an optimizer for the given model.

Parameters:

  • model: The model to optimize
  • optimizer_type: Type of optimizer
  • learning_rate: Learning rate
  • weight_decay: Weight decay coefficient
  • **kwargs: Additional optimizer-specific parameters

Returns:

  • Configured optimizer instance

LionOptimizer

Lion optimizer implementation.

SophiaOptimizer

Sophia optimizer implementation.

get_scheduler

Create a learning rate scheduler.

Parameters:

  • optimizer: Optimizer to schedule
  • scheduler_type: Type of scheduler
  • num_warmup_steps: Number of warmup steps
  • num_training_steps: Total number of training steps
  • **kwargs: Additional scheduler-specific parameters

Returns:

  • Configured learning rate scheduler

CosineWarmupScheduler

Linear warmup followed by cosine decay scheduler.

LinearWarmupScheduler

Linear warmup followed by linear decay scheduler.

GradientClipper

Gradient clipping utility for preventing gradient explosion.

GradientScaler

Gradient scaling utility for mixed precision training.

compress_gradients

Compress gradients for efficient communication in distributed training.

decompress_gradients

Decompress gradients after communication in distributed training.

Evaluation

perplexity

Calculate perplexity metric.

accuracy

Calculate accuracy metric.

f1_score

Calculate F1 score for text generation.

rouge_score

Calculate ROUGE scores (simplified version).

bleu_score

Calculate BLEU score (simplified version).

MetricCalculator

Unified metric calculator for multiple metrics.

Models Module

Architectures

TransformerModel

Transformer model implementation.

Components

MultiHeadAttention

Multi-head attention mechanism.

FeedForward

Feed-forward network.

PositionalEncoding

Positional encoding layer.

LayerNorm

Layer normalization.

Tokenizers Module

BaseTokenizer

Base tokenizer interface.

BPETokenizer

BPE tokenizer implementation.

WordPieceTokenizer

WordPiece tokenizer implementation.

HuggingFaceTokenizer

HuggingFace tokenizer wrapper.

Data Module

get_dataloader

Create a data loader.

TextDataset

Text dataset implementation.

collate_fn

Collate function for data loading.

preprocess_text

Text preprocessing utility.

Inference Module

TextGenerator

Base text generator class.

GreedyGenerator

Greedy text generation.

BeamSearchGenerator

Beam search text generation.

SamplingGenerator

Sampling-based text generation.

generate_text

Generate text using the specified method.

Extensions Module

LoraAdapter

LoRA adapter for model fine-tuning.

AdapterConfig

Configuration for adapter layers.

LoraConfig

Configuration for LoRA adaptation.

apply_lora_to_model

Apply LoRA adaptation to a model.

get_lora_parameters

Get all LoRA parameters from a model.

AcceleratorWrapper

Wrapper for HuggingFace Accelerator.

DistributedTrainer

Distributed trainer using Accelerator.

prepare_model_for_distributed_training

Prepare model for distributed training.

DistributedConfig

Configuration for distributed training.

init_distributed_training

Initialize distributed training environment.

sync_model_parameters

Synchronize model parameters across all processes.

gather_tensors

Gather tensors from all processes.

Utilities Module

setup_logging

Setup logging for the application.

get_logger

Get a logger with the specified name.

log_training_progress

Log training progress.

log_evaluation_results

Log evaluation results.

SystemMonitor

System resource monitor.

ResourceMonitor

Resource usage monitor for training processes.

TrainingMonitor

Training process monitor.

get_system_info

Get comprehensive system information.

check_cuda_availability

Check CUDA availability and get GPU information.

get_available_gpus

Get list of available GPU indices.

estimate_memory_requirements

Estimate memory requirements for training.

select_optimal_device

Select the optimal device for computation.

get_device_memory

Get available memory for a device.

configure_optimal_batch_size

Configure optimal batch size based on available memory.

check_system_resources

Check if system meets minimum resource requirements.

CLI Module

train_command

Main training command.

generate_command

Main generation command.

evaluate_command

Main evaluation command.