From f4aebe4977951d682b53379589ae696dadc3294e Mon Sep 17 00:00:00 2001 From: Hossein Kavianihamedani Date: Mon, 17 Nov 2025 10:22:29 -0800 Subject: [PATCH 1/5] Fix launcher issues --- src/forge/controller/launcher.py | 1 + 1 file changed, 1 insertion(+) diff --git a/src/forge/controller/launcher.py b/src/forge/controller/launcher.py index 2f7e93aec..2a4a470d5 100644 --- a/src/forge/controller/launcher.py +++ b/src/forge/controller/launcher.py @@ -133,6 +133,7 @@ async def initialize(self) -> None: # HostMesh currently requires explicit configuration # of the underlying transport from client to mesh. # This can be removed in the future once this has been removed. + # Changed TCP this to TcpWithHostname for the slurm launcher configure(default_transport=ChannelTransport.TcpWithHostname) async def get_allocator(self, name: str, num_hosts: int) -> tuple[Any, Any, str]: From 88eadcc2c1d1c98aa060dfc3f9de13bc760680df Mon Sep 17 00:00:00 2001 From: Hossein Kavianihamedani Date: Tue, 23 Dec 2025 10:24:41 -0800 Subject: [PATCH 2/5] Refactor: Extract TitanTrainer base class from SFT recipe - Create TitanTrainer dataclass in src/forge/actors/trainer/titan.py - TitanTrainer uses composition with ForgeEngine (self.engine) - Extract common methods: forward_backward, train_step_sft, setup_metric_logger - Update ForgeSFTRecipe to inherit from TitanTrainer - Move compile config to top-level in YAML (required by ForgeJobConfig) - Add documentation for trainer refactoring --- apps/sft/OPTION2_INTEGRATION_GUIDE.md | 172 ++++++++++++++ apps/sft/hf_checkpoint.py | 228 ++++++++++++++++++ apps/sft/llama3_8b.yaml | 22 +- apps/sft/main.py | 296 ++++++++--------------- docs/TRAINER_REFACTORING.md | 322 ++++++++++++++++++++++++++ src/forge/actors/trainer/titan.py | 149 +++++++++++- 6 files changed, 974 insertions(+), 215 deletions(-) create mode 100644 apps/sft/OPTION2_INTEGRATION_GUIDE.md create mode 100644 apps/sft/hf_checkpoint.py create mode 100644 docs/TRAINER_REFACTORING.md diff --git a/apps/sft/OPTION2_INTEGRATION_GUIDE.md b/apps/sft/OPTION2_INTEGRATION_GUIDE.md new file mode 100644 index 000000000..b3e3e6781 --- /dev/null +++ b/apps/sft/OPTION2_INTEGRATION_GUIDE.md @@ -0,0 +1,172 @@ +# Example: How to integrate HuggingFace checkpoint saving into main.py + +## Option 2: Integration Steps + +### Step 1: Import the HF checkpoint module in main.py + +Add to the imports section (around line 13-30): +```python +from forge.apps.sft.hf_checkpoint import save_hf_checkpoint, load_hf_checkpoint +``` + +### Step 2: Replace the existing checkpoint save call + +Find this code in the `train()` method (around line 456): +```python +self.checkpointer.save( + curr_step=self.current_step, + last_step=self.current_step == self.num_training_steps, +) +``` + +Replace it with: +```python +# Save checkpoint using HuggingFace format (bypasses DCP KeyError bug) +checkpoint_config = self.job_config.checkpoint +if (checkpoint_config.enable and + self.current_step % checkpoint_config.interval == 0): + + save_hf_checkpoint( + model=self.model[0], # First model in list + tokenizer=self.tokenizer, + optimizer=self.optimizer.optimizers[0] if self.optimizer else None, + lr_scheduler=self.lr_scheduler.schedulers[0] if self.lr_scheduler else None, + checkpoint_dir=checkpoint_config.folder, + step=self.current_step, + rank=self._rank, + save_optimizer=True, + ) +``` + +### Step 3: Update the config to enable checkpointing + +In `llama3_8b.yaml`, change: +```yaml +checkpoint: + enable: true # ✅ Re-enable checkpointing! + folder: /home/hosseinkh/models/Meta-Llama-3.1-8B-Instruct/checkpoint + initial_load_path: /home/hosseinkh/models/Meta-Llama-3.1-8B-Instruct/ + initial_load_in_hf: true + last_save_in_hf: true # ✅ Now this will work! + interval: 500 + async_mode: "disabled" +``` + +### Step 4: Optional - Load from HF checkpoint at startup + +In the `setup()` method (around line 138), replace: +```python +self.checkpointer.load(step=self.current_step) +``` + +With: +```python +# Load from HuggingFace checkpoint if exists +checkpoint_config = self.job_config.checkpoint +if checkpoint_config.enable and Path(checkpoint_config.folder).exists(): + try: + loaded_step = load_hf_checkpoint( + model=self.model[0], + tokenizer=self.tokenizer, + optimizer=self.optimizer.optimizers[0] if self.optimizer else None, + lr_scheduler=self.lr_scheduler.schedulers[0] if self.lr_scheduler else None, + checkpoint_dir=checkpoint_config.folder, + step=None, # Load latest + rank=self._rank, + ) + self.current_step = loaded_step + logger.info(f"Resumed from step {loaded_step}") + except Exception as e: + logger.warning(f"Could not load checkpoint: {e}, starting from scratch") +``` + +--- + +## How It Works + +### The Problem with PyTorch DCP: +``` +HuggingFace Model (plain state_dict) + ↓ +FSDP Wrapping (changes parameter IDs) + ↓ +Training with new parameter IDs + ↓ +PyTorch DCP Save (tries to map back to original FQNs) + ❌ KeyError: 289 - parameter ID not in fqn_pid_mapping +``` + +### The HuggingFace Solution: +``` +FSDP Model (sharded across GPUs) + ↓ +Gather Full State Dict (FSDP.state_dict_type) + ↓ +Load into Unwrapped Model + ↓ +HuggingFace save_pretrained (safetensors format) + ✅ Works perfectly! +``` + +--- + +## Key Features of HF Checkpoint + +1. **FSDP-Compatible**: Uses `FSDP.state_dict_type()` to properly gather sharded weights +2. **Rank-0 Only**: Only rank 0 saves to avoid race conditions +3. **Barriers**: Proper synchronization across all ranks +4. **Unwrapping**: Recursively unwraps FSDP layers to get original model +5. **SafeTensors**: Uses modern safetensors format +6. **Optimizer State**: Separately saves optimizer/scheduler as PyTorch .pt files +7. **Metadata**: Saves training step and config for easy resumption + +--- + +## File Structure After Saving + +``` +checkpoint_dir/ +├── step_500/ +│ ├── model.safetensors # ✅ Model weights (HF format) +│ ├── config.json # ✅ Model config +│ ├── tokenizer.json # ✅ Tokenizer +│ ├── tokenizer_config.json # ✅ Tokenizer config +│ ├── optimizer.pt # ✅ Optimizer + LR scheduler +│ └── training_metadata.pt # ✅ Step number, etc. +├── step_1000/ +│ └── ... +``` + +--- + +## Benefits Over PyTorch DCP + +| Feature | PyTorch DCP | HuggingFace | +|---------|-------------|-------------| +| FSDP + HF Checkpoint | ❌ KeyError: 289 | ✅ Works | +| Format Compatibility | DCP-specific | ✅ Universal HF format | +| Easy Loading | Complex | ✅ `from_pretrained()` | +| Safetensors Support | Limited | ✅ Native | +| Resume Training | ✅ Yes | ✅ Yes | +| Size | Larger | Smaller (compressed) | + +--- + +## Testing + +After implementing, test with: +```bash +cd /home/hosseinkh/forge_updated/forge +python -m apps.sft.main --config apps/sft/llama3_8b.yaml +``` + +At step 500, you should see: +``` +INFO:forge.apps.sft.hf_checkpoint:Saving HuggingFace checkpoint to /home/hosseinkh/models/Meta-Llama-3.1-8B-Instruct/checkpoint/step_500 +INFO:forge.apps.sft.hf_checkpoint:✅ Model saved to /home/hosseinkh/models/Meta-Llama-3.1-8B-Instruct/checkpoint/step_500 +INFO:forge.apps.sft.hf_checkpoint:✅ Tokenizer saved to /home/hosseinkh/models/Meta-Llama-3.1-8B-Instruct/checkpoint/step_500 +INFO:forge.apps.sft.hf_checkpoint:✅ Optimizer state saved to /home/hosseinkh/models/Meta-Llama-3.1-8B-Instruct/checkpoint/step_500/optimizer.pt +INFO:forge.apps.sft.hf_checkpoint:🎉 Checkpoint saved successfully at step 500 +``` + +No more KeyError! 🎉 diff --git a/apps/sft/hf_checkpoint.py b/apps/sft/hf_checkpoint.py new file mode 100644 index 000000000..77799cf88 --- /dev/null +++ b/apps/sft/hf_checkpoint.py @@ -0,0 +1,228 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +""" +HuggingFace-native checkpoint saving that bypasses PyTorch DCP. + +This module provides checkpoint saving functionality that works with FSDP +by gathering the full model state and using HuggingFace's native save_pretrained. +""" + +import logging +import os +from pathlib import Path +from typing import Optional + +import torch +import torch.distributed as dist +from torch.distributed.fsdp import FullyShardedDataParallel as FSDP +from torch.distributed.fsdp import StateDictType, FullStateDictConfig + +logger = logging.getLogger(__name__) + + +def save_hf_checkpoint( + model: torch.nn.Module, + tokenizer, + optimizer: Optional[torch.optim.Optimizer], + lr_scheduler: Optional[torch.optim.lr_scheduler._LRScheduler], + checkpoint_dir: str, + step: int, + rank: int = 0, + save_optimizer: bool = True, +) -> None: + """ + Save checkpoint in HuggingFace format, bypassing PyTorch DCP. + + This function gathers the full model state from FSDP shards and saves + it using HuggingFace's native format, avoiding the KeyError bug in DCP. + + Args: + model: The model to save (can be FSDP-wrapped) + tokenizer: The tokenizer to save + optimizer: Optional optimizer to save + lr_scheduler: Optional LR scheduler to save + checkpoint_dir: Base directory for checkpoints + step: Current training step + rank: Process rank (only rank 0 saves) + save_optimizer: Whether to save optimizer state + """ + # Create checkpoint directory + ckpt_path = Path(checkpoint_dir) / f"step_{step}" + + if rank == 0: + ckpt_path.mkdir(parents=True, exist_ok=True) + logger.info(f"Saving HuggingFace checkpoint to {ckpt_path}") + + # Wait for directory creation + if dist.is_initialized(): + dist.barrier() + + # Handle FSDP model + if isinstance(model, FSDP): + # Configure FSDP to gather full state dict + save_policy = FullStateDictConfig(offload_to_cpu=True, rank0_only=True) + with FSDP.state_dict_type( + model, + StateDictType.FULL_STATE_DICT, + state_dict_config=save_policy, + ): + model_state = model.state_dict() + else: + model_state = model.state_dict() + + # Only rank 0 saves to avoid race conditions + if rank == 0: + # Get the unwrapped model for HF saving + if isinstance(model, FSDP): + # Access the original model inside FSDP wrapper + unwrapped_model = model._fsdp_wrapped_module + # If it's still wrapped (nested FSDP), keep unwrapping + while hasattr(unwrapped_model, '_fsdp_wrapped_module'): + unwrapped_model = unwrapped_model._fsdp_wrapped_module + else: + unwrapped_model = model + + # Load the full state dict into unwrapped model + unwrapped_model.load_state_dict(model_state, strict=True) + + # Save using HuggingFace's native method + try: + unwrapped_model.save_pretrained( + ckpt_path, + safe_serialization=True, # Use safetensors format + ) + logger.info(f"✅ Model saved to {ckpt_path}") + except Exception as e: + logger.error(f"❌ Failed to save model: {e}") + raise + + # Save tokenizer + if tokenizer is not None: + try: + tokenizer.save_pretrained(ckpt_path) + logger.info(f"✅ Tokenizer saved to {ckpt_path}") + except Exception as e: + logger.warning(f"⚠️ Failed to save tokenizer: {e}") + + # Save optimizer and scheduler (as PyTorch native - no HF format for these) + if save_optimizer and optimizer is not None: + try: + opt_path = ckpt_path / "optimizer.pt" + optimizer_state = { + 'optimizer': optimizer.state_dict(), + 'step': step, + } + if lr_scheduler is not None: + optimizer_state['lr_scheduler'] = lr_scheduler.state_dict() + + torch.save(optimizer_state, opt_path) + logger.info(f"✅ Optimizer state saved to {opt_path}") + except Exception as e: + logger.warning(f"⚠️ Failed to save optimizer: {e}") + + # Save training metadata + try: + meta_path = ckpt_path / "training_metadata.pt" + metadata = { + 'step': step, + 'model_config': unwrapped_model.config.to_dict() if hasattr(unwrapped_model, 'config') else {}, + } + torch.save(metadata, meta_path) + logger.info(f"✅ Metadata saved to {meta_path}") + except Exception as e: + logger.warning(f"⚠️ Failed to save metadata: {e}") + + logger.info(f"🎉 Checkpoint saved successfully at step {step}") + + # Synchronize all ranks + if dist.is_initialized(): + dist.barrier() + + +def load_hf_checkpoint( + model: torch.nn.Module, + tokenizer, + optimizer: Optional[torch.optim.Optimizer], + lr_scheduler: Optional[torch.optim.lr_scheduler._LRScheduler], + checkpoint_dir: str, + step: Optional[int] = None, + rank: int = 0, +) -> int: + """ + Load checkpoint from HuggingFace format. + + Args: + model: The model to load into + tokenizer: The tokenizer (not modified, HF loads from config) + optimizer: Optional optimizer to load state into + lr_scheduler: Optional LR scheduler to load state into + checkpoint_dir: Base directory for checkpoints + step: Specific step to load (None = latest) + rank: Process rank + + Returns: + The step number that was loaded + """ + ckpt_base = Path(checkpoint_dir) + + # Find checkpoint to load + if step is None: + # Find latest checkpoint + checkpoints = sorted([d for d in ckpt_base.glob("step_*") if d.is_dir()]) + if not checkpoints: + logger.warning(f"No checkpoints found in {checkpoint_dir}") + return 0 + ckpt_path = checkpoints[-1] + step = int(ckpt_path.name.split("_")[1]) + else: + ckpt_path = ckpt_base / f"step_{step}" + if not ckpt_path.exists(): + raise FileNotFoundError(f"Checkpoint not found: {ckpt_path}") + + if rank == 0: + logger.info(f"Loading HuggingFace checkpoint from {ckpt_path}") + + # Get unwrapped model + if isinstance(model, FSDP): + unwrapped_model = model._fsdp_wrapped_module + while hasattr(unwrapped_model, '_fsdp_wrapped_module'): + unwrapped_model = unwrapped_model._fsdp_wrapped_module + else: + unwrapped_model = model + + # Load model using HuggingFace's from_pretrained + try: + from transformers import AutoModelForCausalLM + loaded_model = AutoModelForCausalLM.from_pretrained( + ckpt_path, + torch_dtype=unwrapped_model.dtype if hasattr(unwrapped_model, 'dtype') else torch.float32, + ) + unwrapped_model.load_state_dict(loaded_model.state_dict(), strict=True) + if rank == 0: + logger.info(f"✅ Model loaded from {ckpt_path}") + except Exception as e: + logger.error(f"❌ Failed to load model: {e}") + raise + + # Load optimizer and scheduler + if optimizer is not None: + opt_path = ckpt_path / "optimizer.pt" + if opt_path.exists(): + try: + optimizer_state = torch.load(opt_path, map_location='cpu') + optimizer.load_state_dict(optimizer_state['optimizer']) + if lr_scheduler is not None and 'lr_scheduler' in optimizer_state: + lr_scheduler.load_state_dict(optimizer_state['lr_scheduler']) + if rank == 0: + logger.info(f"✅ Optimizer state loaded from {opt_path}") + except Exception as e: + logger.warning(f"⚠️ Failed to load optimizer: {e}") + + if rank == 0: + logger.info(f"🎉 Checkpoint loaded successfully from step {step}") + + return step diff --git a/apps/sft/llama3_8b.yaml b/apps/sft/llama3_8b.yaml index d9a5a9783..669aaf2df 100644 --- a/apps/sft/llama3_8b.yaml +++ b/apps/sft/llama3_8b.yaml @@ -12,7 +12,7 @@ model_name: "meta-llama/Meta-Llama-3.1-8B-Instruct" model: name: llama3 flavor: 8B - hf_assets_path: hf://${model_name} + hf_assets_path: /home/hosseinkh/models/Meta-Llama-3.1-8B-Instruct processes: procs: 8 @@ -26,18 +26,20 @@ optimizer: lr_scheduler: warmup_steps: 200 +compile: + enable: false + training: local_batch_size: 8 seq_len: 2048 max_norm: 1.0 - steps: 1000 - compile: false + steps: 300 datasets: - path: "yahma/alpaca-cleaned" split: "train[:95%]" eval: - eval_every_n_steps: 50 # null = disabled + eval_every_n_steps: 10 # null = disabled max_eval_steps: null # null = run until epoch completes datasets: - path: "yahma/alpaca-cleaned" @@ -54,22 +56,18 @@ parallelism: checkpoint: enable: true - folder: ./checkpoint # The folder to save checkpoints to. - initial_load_path: hf://${model_name} # The path to load the initial checkpoint from. Ignored if `folder` exists. + folder: /home/hosseinkh/models/Meta-Llama-3.1-8B-Instruct/checkpoint # The folder to save checkpoints to. + initial_load_path: /home/hosseinkh/models/Meta-Llama-3.1-8B-Instruct # The path to load the initial checkpoint from. Ignored if `folder` exists. initial_load_in_hf: true # If true, interpret initial_load_path as a HuggingFace model repo last_save_in_hf: true - interval: 500 + interval: 50 async_mode: "disabled" activation_checkpoint: mode: selective selective_ac_option: op -metric_logging: - wandb: - project: sft-training - group: sft_exp_${oc.env:USER} - logging_mode: global_reduce # global_reduce, per_rank_reduce, per_rank_no_reduce +metric_logging: {} # profiling: diff --git a/apps/sft/main.py b/apps/sft/main.py index 4f2a7be74..2056ab3d8 100644 --- a/apps/sft/main.py +++ b/apps/sft/main.py @@ -20,8 +20,7 @@ import torch -import torchtitan.experiments.forge.train_spec as forge_train_spec -from forge.controller import ForgeActor +from forge.actors.trainer import TitanTrainer from forge.data.collate import collate_padded from forge.data.datasets.sft_dataset import AlpacaToMessages, sft_iterable_dataset from forge.data.tokenizer import HuggingFaceModelTokenizer @@ -31,17 +30,9 @@ from monarch.actor import current_rank, current_size, endpoint from omegaconf import DictConfig, OmegaConf -from torch import nn from torchdata.stateful_dataloader import StatefulDataLoader -from torchtitan.components.loss import LossFunction -from torchtitan.components.lr_scheduler import LRSchedulersContainer -from torchtitan.components.optimizer import OptimizersContainer -from torchtitan.distributed import ParallelDims, utils as dist_utils -from torchtitan.experiments.forge.engine import ForgeEngine from torchtitan.experiments.forge.job_config import ForgeJobConfig -# from tqdm import tqdm - # stubs for now Checkpointer = Any Dataloader = Any @@ -53,62 +44,83 @@ logger.setLevel(logging.INFO) -class ForgeSFTRecipe(ForgeActor, ForgeEngine): - job_config: ForgeJobConfig - train_spec: forge_train_spec.ForgeTrainSpec - parallel_dims: ParallelDims - model: list[nn.Module] - loss_fn: LossFunction - optimizer: OptimizersContainer - lr_scheduler: LRSchedulersContainer - checkpointer: Checkpointer - tokenizer: Tokenizer - train_dataloader: Dataloader - # val_dataloader: Dataloader - metric_logger: MetricLogger - profiler: Profiler - device: torch.device - step: int +class ForgeSFTRecipe(TitanTrainer): + """SFT Recipe built on TitanTrainer. + + Inherits from TitanTrainer which provides: + - rank_should_record_loss: Only last PP stage records loss + - record_batch_metrics(): Generic batch metric recording + - setup_metric_logger(): Metric logger initialization + - forward_backward(): Forward/backward with PP+CP support + - train_step_sft(): SFT-style training step + + Uses composition pattern: access engine via self.engine.X + + This class adds SFT-specific functionality: + - Data loading (tokenizer, datasets) + - Evaluation loop + - Training loop with eval triggers + """ def __init__(self, config: DictConfig): - job_config = ForgeJobConfig().to_dict() - # Hack to deal with literal types from titan - job_config = OmegaConf.merge(job_config, config) + # Get valid ForgeJobConfig fields + forge_job_config_defaults = ForgeJobConfig().to_dict() + + # Store the full config for SFT-specific fields + full_config = OmegaConf.to_container(config, resolve=True) + + # Only merge config keys that are valid ForgeJobConfig fields + # Filter out non-ForgeJobConfig fields like model_name, processes, metric_logging, eval, etc. + filtered_config = { + k: v for k, v in full_config.items() if k in forge_job_config_defaults + } + + # Store datasets separately before filtering them out (they're SFT-specific, not in Training dataclass) + training_datasets = None + if "training" in filtered_config and "datasets" in filtered_config["training"]: + training_datasets = filtered_config["training"].pop("datasets") + + # Store eval config separately (SFT-specific, not in ForgeJobConfig) + self.eval_config = full_config.get("eval", {}) + + job_config = OmegaConf.merge(forge_job_config_defaults, filtered_config) + + self.job_config = ForgeJobConfig(**job_config) + # Restore datasets to job_config for SFT use + if training_datasets is not None: + self.job_config.training.datasets = training_datasets self.current_step = 0 - self.num_training_steps = job_config.training.steps - self.gradient_accumulation_steps = 1 # Example value, adjust as needed + self.num_training_steps = self.job_config.training.steps + self.gradient_accumulation_steps = 1 self._rank = current_rank().rank self._size = math.prod(current_size().values()) - super().__init__(job_config) - async def setup_metric_logger(self): - """Initialization happens in the main process. Here we just retrieve it""" - mlogger = await get_or_create_metric_logger() - return mlogger + # Convert job_config to dict for TitanTrainer, ensuring datasets is removed + titan_config = OmegaConf.to_container(job_config, resolve=True) + if "training" in titan_config and "datasets" in titan_config["training"]: + del titan_config["training"]["datasets"] + + # Initialize TitanTrainer with job config fields (without datasets) + super().__init__(**titan_config) - def record_batch_metrics(self, data_metrics: list): - """Since the dataloader creates new processes, we dont call `record_metric` in the dataset. - Instead, pop the metrics from the batch and record them here.""" - for metric in data_metrics: - record_metric(metric.key, metric.value, metric.reduction) + # setup_metric_logger() inherited from TitanTrainer + # record_batch_metrics() inherited from TitanTrainer @endpoint async def setup(self): + # Call parent's setup helper to initialize engine and rank_should_record_loss + # self._setup_engine() + # Validate that compile is only used with flex attention - if self.job_config.training.compile: + if self.job_config.compile.enable: raise ValueError( - "training.compile=True is not currently supported. " + "compile.enable=True is not currently supported. " "Compile is only supported with flex attention enabled, which requires PyTorch nightly. " - "Please set training.compile=false in your config." + "Please set compile.enable=false in your config." ) - # all ranks should record loss, except when PP=True. Then, only the last stage should record loss. - self.rank_should_record_loss = True - if hasattr(self, "pp_has_last_stage") and not self.pp_has_last_stage: - self.rank_should_record_loss = False - - # metric logger + # metric logger (inherited from TitanTrainer) self.mlogger = await self.setup_metric_logger() # Load training datasets @@ -116,11 +128,10 @@ async def setup(self): train_datasets_config = self.job_config.training.datasets self.train_dataloader = self.setup_data(train_datasets_config) - # Load eval datasets - eval_config = self.job_config["eval"] + # Load eval datasets (using self.eval_config stored in __init__) self.val_dataloaders = {} - self.eval_every_n_steps = eval_config["eval_every_n_steps"] - max_eval_steps = eval_config["max_eval_steps"] + self.eval_every_n_steps = self.eval_config.get("eval_every_n_steps") + max_eval_steps = self.eval_config.get("max_eval_steps") self.max_eval_steps = ( max_eval_steps if max_eval_steps and max_eval_steps > 0 else None ) @@ -129,7 +140,7 @@ async def setup(self): ) if self.validation_enabled: logger.info("Setting eval datasets") - self.eval_datasets_config = eval_config.datasets + self.eval_datasets_config = self.eval_config.get("datasets", []) for i, dataset_config in enumerate(self.eval_datasets_config): ds_name = dataset_config.get("dataset_name", i) @@ -140,22 +151,16 @@ async def setup(self): # TODO: confirm that this is working properly # Should also use load, not dcp_load - self.checkpointer.load(step=self.current_step) - - # self.profiler = self.setup_profiler(self.train_config.profiler_config) - # self.logger = self.setup_logger(self.train_config.logger_config) + self.engine.checkpointer.load(step=self.current_step) def setup_data(self, dataset_configs: list[dict]) -> StatefulDataLoader: """Instantiates datasets and returns a StatefulDataLoader. Args: - dataset_configs (list[dict]): List of dataset config dicts used as `sft_iterable_dataset(**dataset_configs[i])`. + dataset_configs (list[dict]): List of dataset config dicts. Returns: StatefulDataLoader - - Raises: - ValueError: If multiple datasets provided (not yet supported) """ # TODO felipemello: Currently only support single dataset @@ -167,7 +172,6 @@ def setup_data(self, dataset_configs: list[dict]) -> StatefulDataLoader: dataset_config = dataset_configs[0] - # TODO: Evaluate if tokenizers should be created once and shared for every dataset # Load tokenizer tokenizer = HuggingFaceModelTokenizer( tokenizer_json_path=os.path.join( @@ -190,10 +194,13 @@ def setup_data(self, dataset_configs: list[dict]) -> StatefulDataLoader: ), ) - # Get DP mesh for data sharding + # Get DP mesh for data sharding (use self.engine for composition) dp_mesh = None - if self.parallel_dims is not None and self.parallel_dims.dp_enabled: - dp_mesh = self.parallel_dims.world_mesh.get_group("dp") + if ( + self.engine.parallel_dims is not None + and self.engine.parallel_dims.dp_enabled + ): + dp_mesh = self.engine.parallel_dims.world_mesh.get_group("dp") # Pass config directly to dataset constructor dataset = sft_iterable_dataset( @@ -211,118 +218,28 @@ def setup_data(self, dataset_configs: list[dict]) -> StatefulDataLoader: return dataloader - def forward_backward( - self, - input_dict: dict[str, torch.Tensor], - labels: torch.Tensor, - skip_backward: bool = False, - ) -> torch.Tensor: - model_parts = self.model_parts - parallel_dims = self.parallel_dims - - # apply context parallelism if cp is enabled - # ensure CP handles the separate freqs_cis buffer for each pp stage - inputs = input_dict["tokens"] - optional_context_parallel_ctx = ( - dist_utils.create_context_parallel_ctx( - cp_mesh=parallel_dims.world_mesh["cp"], - cp_buffers=[inputs, labels] + [m.freqs_cis for m in model_parts], - cp_seq_dims=[1, 1] + [0 for _ in model_parts], - cp_no_restore_buffers={inputs, labels}, - cp_rotate_method=self.job_config.parallelism.context_parallel_rotate_method, - ) - if parallel_dims.cp_enabled - else None - ) - - if parallel_dims.pp_enabled: - # Pipeline Parallel forward / backward inside step() call - with self.train_context(optional_context_parallel_ctx): - targets, losses = ( - (labels, []) if self.pp_has_last_stage else (None, None) - ) - if self.pp_has_first_stage: - self.pp_schedule.step(inputs, target=targets, losses=losses) - else: - self.pp_schedule.step(target=targets, losses=losses) - - # accumulate losses across pipeline microbatches - # TODO: PP+FSDP unexpectedly puts the loss back to the CPU - loss = ( - torch.sum(torch.stack(losses)).to(self.device) - if self.pp_has_last_stage - else torch.tensor(-1.0, device=self.device) - ) - - # TODO: PP requires gradients enabled and cant deactive with no_grad - if skip_backward: - loss = loss.detach() - - else: - # Non-PP forward / backward - with self.train_context(optional_context_parallel_ctx): - assert len(model_parts) == 1 - with self.maybe_enable_amp: - pred = model_parts[0](inputs) - loss = self.loss_fn(pred, labels) - # need to free to before bwd to avoid peaking memory - del pred - - # Only run backward if requested. Useful for eval. - if not skip_backward: - loss.backward() - - return loss - - def train_step(self, batch) -> None: - # TODO - # with GradientAccumulation( - # self.gradient_accumulation_steps, - # self.model, - # self.data_parallel_size, - # ) as grad_acc: - labels = batch.pop("labels") - loss = self.forward_backward(batch, labels) - - if self.rank_should_record_loss: - loss_val = loss.item() - record_metric("ForgeSFTRecipe/train_step/loss", loss_val, Reduce.MEAN) - logger.info( - f"step {self.current_step} / {self.num_training_steps} | Loss: {loss_val}" - ) - - # self.pbar.set_description(f"{self.current_step}|Loss: {loss}") - # self.pbar.update(1) - self.optimizers.step() - self.lr_schedulers.step() + # forward_backward() inherited from TitanTrainer + # train_step_sft() inherited from TitanTrainer async def evaluate(self) -> None: - """Run evaluation on multiple datasets, one at a time. - - 1. Set models to eval mode - 2. For each eval dataset: - - Create fresh iterator (starts from epoch 0) - - Use StopAfterOneEpoch to iterate until epoch boundary. This utility - is necessary for infinite iterable dataset, since epoch boundaries are not known. - - Respect max_eval_steps cap if configured - - Record loss and step metrics (on dp rank only) - 3. Restore models to train mode - """ + """Run evaluation on multiple datasets, one at a time.""" - # Set models to eval mode - for model_part in self.model_parts: + # Set models to eval mode (use self.engine for composition) + for model_part in self.engine.model_parts: model_part.eval() # Get DP process group for epoch synchronization dp_mesh = None - if self.parallel_dims is not None and self.parallel_dims.dp_enabled: - dp_mesh = self.parallel_dims.world_mesh.get_group("dp") + if ( + self.engine.parallel_dims is not None + and self.engine.parallel_dims.dp_enabled + ): + dp_mesh = self.engine.parallel_dims.world_mesh.get_group("dp") # For non-PP: disable gradients to save memory - # TODO: For PP, if disabling gradients, throws error maybe_no_grad = ( contextlib.nullcontext() - if self.parallel_dims.pp_enabled + if self.engine.parallel_dims.pp_enabled else torch.no_grad() ) @@ -333,19 +250,17 @@ async def evaluate(self) -> None: logger.info(f"=====Evaluating dataset: {dataset_name}=====") # Evaluation loop for this dataset - total_loss = torch.tensor(0.0, device=self.device) + total_loss = torch.tensor(0.0, device=self.engine.device) num_steps = 0 - # NOTE: Assumes batch contains field "metrics" containing "num_epochs" batch_iter = StopAfterOneEpoch( - iter=iter(val_dataloader), # Fresh iterator from epoch 0, - device=self.device, + iter=iter(val_dataloader), + device=self.engine.device, dp_mesh=dp_mesh, ) with maybe_no_grad: for batch in batch_iter: - # if max_eval_steps>len(dataset), it will be stopped earlier by StopAfterOneEpoch. if ( self.max_eval_steps is not None and num_steps >= self.max_eval_steps @@ -358,22 +273,20 @@ async def evaluate(self) -> None: # Move tensors to device for key, value in batch.items(): if isinstance(value, torch.Tensor): - batch[key] = value.to(self.device) + batch[key] = value.to(self.engine.device) - # Process batch + # Process batch - use inherited forward_backward labels = batch.pop("labels") loss = self.forward_backward(batch, labels, skip_backward=True) total_loss += loss num_steps += 1 - # Log progress if self.rank_should_record_loss: loss_val = loss.item() logger.info( f"[dataset {dataset_name}] Step {num_steps} | Loss: {loss_val:.4f}" ) - # log loss avg_loss = (total_loss / max(num_steps, 1)).item() all_dataset_losses.append(avg_loss) all_dataset_steps.append(num_steps) @@ -387,13 +300,11 @@ async def evaluate(self) -> None: Reduce.MEAN, ) - # Record macro and micro average losses across datasets (only if multiple datasets) + # Record macro and micro average losses across datasets if self.rank_should_record_loss and len(all_dataset_losses) > 1: - # Macro: same weight for all datasets macro_avg_loss = sum(all_dataset_losses) / len(all_dataset_losses) record_metric("evaluate/macro_avg_loss", macro_avg_loss, Reduce.MEAN) - # Micro: weighted mean by dataset size total_steps = sum(all_dataset_steps) micro_avg_loss = ( sum( @@ -410,7 +321,7 @@ async def evaluate(self) -> None: ) # Restore train mode - for model_part in self.model_parts: + for model_part in self.engine.model_parts: model_part.train() logger.info("==Evaluation complete==") @@ -418,25 +329,22 @@ async def evaluate(self) -> None: @endpoint async def train(self) -> None: dataloader = iter(self.train_dataloader) - self.optimizers.zero_grad() - - # TODO: tqdm is broken in Monarch actors - # self.pbar = tqdm(initial=self.current_step, total=self.num_training_steps) + self.engine.optimizers.zero_grad() while self.current_step < self.num_training_steps: batch = next(dataloader) - # Pop and record metrics from batch before moving to device + # Pop and record metrics from batch (inherited from TitanTrainer) self.record_batch_metrics(batch.pop("metrics", [])) record_metric("ForgeSFTRecipe/train/step", self.current_step, Reduce.MEAN) # Move tensors to the appropriate device for k, v in batch.items(): if isinstance(v, torch.Tensor): - batch[k] = v.to("cuda") # TODO: hardcoded for now + batch[k] = v.to(self.engine.device) - self.train_step(batch) - # self.profiler.step() + # Use inherited train_step_sft + self.train_step_sft(batch) self.current_step += 1 # Run evaluation periodically if enabled @@ -446,7 +354,7 @@ async def train(self) -> None: ): await self.evaluate() - self.checkpointer.save( + self.engine.checkpointer.save( curr_step=self.current_step, last_step=self.current_step == self.num_training_steps, ) @@ -455,21 +363,19 @@ async def train(self) -> None: if self._rank == 0: await self.mlogger.flush.call_one(global_step=self.current_step) - # self.pbar.close() - if self.validation_enabled: logger.info("Running final evaluation at end of training...") await self.evaluate() @endpoint async def cleanup(self) -> None: - if self.checkpointer: - self.checkpointer.close() + if self.engine.checkpointer: + self.engine.checkpointer.close() if getattr(self, "mlogger", None): await self.mlogger.shutdown.call_one() def __repr__(self) -> str: - return "Trainer" + return "ForgeSFTRecipe" async def run(cfg: DictConfig) -> None: diff --git a/docs/TRAINER_REFACTORING.md b/docs/TRAINER_REFACTORING.md new file mode 100644 index 000000000..3bdc02ec5 --- /dev/null +++ b/docs/TRAINER_REFACTORING.md @@ -0,0 +1,322 @@ +# Trainer Refactoring RFC: Unified TitanTrainer + +## Overview + +This RFC defines the target class hierarchy for unifying SFT and RL trainers. + +--- + +## Target Class Hierarchy + +``` +ForgeActor (base actor) + │ + └── TitanTrainer (generic trainer with ForgeEngine composition) + │ + └── ForgeSFTRecipe (SFT-specific: data loading, eval, training loop) + +ForgeEngine (standalone training engine - used via composition, not inheritance) +``` + +--- + +## TitanTrainer: The Base Trainer + +```python +@dataclass +class TitanTrainer(ForgeActor): + """Base trainer that wraps ForgeEngine with composition pattern. + + Provides reusable components for any training workload: + - _setup_engine(): Engine initialization + - rank_should_record_loss: Only last PP stage records loss + - record_batch_metrics(): Generic batch metric recording + - setup_metric_logger(): Metric logger initialization + - forward_backward(): Forward/backward with PP+CP support (uses engine.loss_fn) + - forward_backward_rl(): Forward/backward for RL (uses custom self.loss) + - train_step_sft(): SFT-style training step + - train_step(): RL-style training step (endpoint) + - push_weights(): Push weights to torchstore + - cleanup(): Cleanup resources + + Uses COMPOSITION: access ForgeEngine via self.engine.X + """ + + # === ForgeJobConfig fields === + job: Job = field(default_factory=Job) + model: Model = field(default_factory=Model) + optimizer: Optimizer = field(default_factory=Optimizer) + lr_scheduler: LRScheduler = field(default_factory=LRScheduler) + training: Training = field(default_factory=Training) + parallelism: Parallelism = field(default_factory=Parallelism) + checkpoint: Checkpoint = field(default_factory=Checkpoint) + activation_checkpoint: ActivationCheckpoint = field(default_factory=ActivationCheckpoint) + compile: Compile = field(default_factory=Compile) + quantize: Quantize = field(default_factory=Quantize) + comm: Comm = field(default_factory=Comm) + memory_estimation: MemoryEstimation = field(default_factory=MemoryEstimation) + + # === Non-JobConfig fields === + loss: Callable = lambda logits, **targets: logits # Custom loss for RL + + # === Engine (set during setup) === + engine: ForgeEngine # Composition - not inheritance! + rank_should_record_loss: bool + + # === Lifecycle Methods === + + @endpoint + async def setup(self): + self._setup_engine() + + def _setup_engine(self): + """Initialize ForgeEngine. Non-endpoint helper for subclasses.""" + self.engine = ForgeEngine(ForgeJobConfig(...)) + # Set rank_should_record_loss based on PP stage + self.rank_should_record_loss = True + if hasattr(self.engine, "pp_has_last_stage") and not self.engine.pp_has_last_stage: + self.rank_should_record_loss = False + + @endpoint + async def cleanup(self) -> None: + if self.engine.checkpointer: + self.engine.checkpointer.close() + + # === Reusable Utilities === + + async def setup_metric_logger(self): + """Get or create metric logger.""" + return await get_or_create_metric_logger() + + def record_batch_metrics(self, data_metrics: list): + """Record metrics from batch.""" + for metric in data_metrics: + record_metric(metric.key, metric.value, metric.reduction) + + # === Forward/Backward Passes === + + def forward_backward( + self, + input_dict: dict[str, Tensor], + labels: Tensor, + skip_backward: bool = False + ) -> Tensor: + """Forward/backward for SFT with PP+CP support. + + Uses engine.loss_fn (built-in cross-entropy loss). + Supports Pipeline Parallelism and Context Parallelism. + """ + # ... PP and CP handling ... + loss = self.engine.loss_fn(pred, labels) + if not skip_backward: + loss.backward() + return loss + + def forward_backward_rl( + self, + inputs: dict[str, Tensor], + targets: dict[str, Tensor] + ) -> Tensor: + """Forward/backward for RL with custom loss. + + Uses self.loss (custom RL loss function passed to trainer). + """ + logits = model(**inputs) + loss = self.loss(logits, **targets) # Custom loss! + loss.backward() + return loss + + # === Training Steps === + + def train_step_sft(self, batch: dict[str, Tensor]) -> None: + """SFT training step. Called from internal training loop. + + Extracts labels, calls forward_backward, logs metrics, optimizer step. + """ + labels = batch.pop("labels") + loss = self.forward_backward(batch, labels) + + if self.rank_should_record_loss: + record_metric("loss", loss.item(), Reduce.MEAN) + + self.engine.optimizers.step() + self.engine.lr_schedulers.step() + + @endpoint + async def train_step( + self, + inputs: list[dict[str, Tensor]], + targets: list[dict[str, Tensor]] + ) -> float: + """RL training step endpoint. Called from external control loop. + + For GRPO and similar algorithms where data comes from external rollouts. + """ + loss = self.forward_backward_rl(inputs[self.engine.dp_rank], targets[self.engine.dp_rank]) + torch.distributed.all_reduce(loss) + + self.engine.optimizers.step() + self.engine.optimizers.zero_grad() + self.engine.lr_schedulers.step() + + return loss.item() +``` + +--- + +## ForgeSFTRecipe: SFT-Specific Recipe + +```python +class ForgeSFTRecipe(TitanTrainer): + """SFT Recipe that adds data loading, eval, and training loop. + + INHERITS from TitanTrainer: + ✓ _setup_engine(), rank_should_record_loss + ✓ setup_metric_logger(), record_batch_metrics() + ✓ forward_backward(), train_step_sft() + + ADDS SFT-specific: + + __init__: Config handling (datasets, eval config) + + setup_data(): Create StatefulDataLoader with tokenizer + + evaluate(): Evaluation loop + + train(): Training loop with checkpointing and eval triggers + """ + + def __init__(self, config: DictConfig): + # Extract SFT-specific config (datasets, eval) + self.eval_config = config.get("eval", {}) + training_datasets = config.training.pop("datasets", None) + + # Initialize parent with ForgeJobConfig fields only + super().__init__(**titan_config) + + # Restore datasets for SFT use + self.job_config.training.datasets = training_datasets + + @endpoint + async def setup(self): + # Initialize engine (from TitanTrainer) + self._setup_engine() + + # Initialize metric logger (from TitanTrainer) + self.mlogger = await self.setup_metric_logger() + + # SFT-specific: Setup data loaders + self.train_dataloader = self.setup_data(self.job_config.training.datasets) + + # SFT-specific: Setup eval dataloaders + if self.validation_enabled: + for dataset_config in self.eval_config.datasets: + self.val_dataloaders[name] = self.setup_data([dataset_config]) + + def setup_data(self, dataset_configs: list[dict]) -> StatefulDataLoader: + """SFT-specific: Create dataloader with tokenizer.""" + tokenizer = HuggingFaceModelTokenizer(...) + dataset = sft_iterable_dataset(tokenizer, ...) + return StatefulDataLoader(dataset, ...) + + async def evaluate(self) -> None: + """SFT-specific: Evaluation loop.""" + for model_part in self.engine.model_parts: + model_part.eval() + + for batch in val_dataloader: + labels = batch.pop("labels") + # Use inherited forward_backward with skip_backward=True + loss = self.forward_backward(batch, labels, skip_backward=True) + + for model_part in self.engine.model_parts: + model_part.train() + + @endpoint + async def train(self) -> None: + """SFT-specific: Training loop with eval triggers.""" + self.engine.optimizers.zero_grad() + + while self.current_step < self.num_training_steps: + batch = next(self.train_dataloader) + + # Use inherited utilities + self.record_batch_metrics(batch.pop("metrics", [])) + + # Move to device + batch = {k: v.to(self.engine.device) for k, v in batch.items()} + + # Use inherited train_step_sft + self.train_step_sft(batch) + self.current_step += 1 + + # Eval periodically + if self.current_step % self.eval_every_n_steps == 0: + await self.evaluate() + + # Checkpoint + self.engine.checkpointer.save(curr_step=self.current_step, ...) +``` + +--- + +## GRPO Usage (External Control Loop) + +```python +# GRPO uses TitanTrainer directly (no subclass needed for now) +# Data/rollouts managed by external Orchestrator + +async def run_grpo(): + # Create trainer with custom GRPO loss + trainer = await TitanTrainer.options(...).as_actor( + loss=grpo_loss_fn, # Custom RL loss + model={"name": "llama3", ...}, + ) + await trainer.setup.call() + + # External control loop (Orchestrator) + for epoch in range(num_epochs): + # Generate rollouts (external) + rollouts = await generate_rollouts(...) + + # Train on rollouts using train_step endpoint + for batch in rollouts: + inputs = [{"input_ids": batch.ids, "attention_mask": batch.mask}] + targets = [{"advantages": batch.advs, "old_log_probs": batch.log_probs}] + loss = await trainer.train_step.call(inputs, targets) + + await trainer.cleanup.call() +``` + +--- + +## Key Design Decisions + +| Decision | Rationale | +|----------|-----------| +| **Composition for ForgeEngine** | TitanTrainer wraps engine via `self.engine.X` - avoids diamond inheritance | +| **Inheritance for Recipes** | ForgeSFTRecipe inherits from TitanTrainer - gets all utilities automatically | +| **Two forward_backward variants** | `forward_backward()` uses engine.loss_fn (SFT); `forward_backward_rl()` uses custom self.loss (RL) | +| **Endpoint vs Non-endpoint** | `train_step()` is endpoint for external control (RL); `train_step_sft()` is regular method for internal loop (SFT) | +| **`_setup_engine()` helper** | Non-endpoint method so subclasses can call it from their own setup endpoints | + +--- + +## What's Reusable in TitanTrainer + +| Component | SFT Uses | RL Uses | Notes | +|-----------|----------|---------|-------| +| `_setup_engine()` | ✓ | ✓ | Initialize ForgeEngine | +| `rank_should_record_loss` | ✓ | ✓ | PP stage handling | +| `setup_metric_logger()` | ✓ | ✓ | Metric logging | +| `record_batch_metrics()` | ✓ | - | Batch metrics from dataloader | +| `forward_backward()` | ✓ | - | Uses engine.loss_fn | +| `forward_backward_rl()` | - | ✓ | Uses custom self.loss | +| `train_step_sft()` | ✓ | - | Internal training step | +| `train_step()` endpoint | - | ✓ | External training step | +| `push_weights()` | - | ✓ | Push to torchstore | +| `cleanup()` | ✓ | ✓ | Resource cleanup | + +--- + +## Migration Path + +1. **PR 1 (Current)**: Add reusable components to TitanTrainer, have ForgeSFTRecipe inherit from TitanTrainer +2. **PR 2 (Future)**: Migrate GRPO to use TitanTrainer if beneficial +3. **PR 3 (Future)**: Add more reusable components (gradient clipping, etc.) diff --git a/src/forge/actors/trainer/titan.py b/src/forge/actors/trainer/titan.py index 6a4fa3cb4..adcaebdec 100644 --- a/src/forge/actors/trainer/titan.py +++ b/src/forge/actors/trainer/titan.py @@ -6,7 +6,6 @@ import logging import os - import time from collections.abc import Mapping from dataclasses import dataclass, field, fields @@ -25,6 +24,7 @@ from forge.controller import ForgeActor from forge.data.utils import batch_to_device +from forge.observability.metric_actors import get_or_create_metric_logger from forge.observability.metrics import record_metric, Reduce from forge.observability.perf_tracker import Tracer @@ -45,6 +45,7 @@ Quantize, Training, ) +from torchtitan.distributed import utils as dist_utils from torchtitan.experiments.forge.engine import ForgeEngine from torchtitan.experiments.forge.job_config import ForgeJobConfig @@ -72,6 +73,13 @@ class TitanTrainer(ForgeActor): The trainer handles: - Forward and backward propagation with automatic mixed precision (AMP) - Optimizer steps with learning rate scheduling + + Provides reusable components for SFT and RL workloads: + - rank_should_record_loss: Only last PP stage records loss + - record_batch_metrics(): Generic batch metric recording + - setup_metric_logger(): Metric logger initialization + - forward_backward(): Forward/backward with PP+CP support + - train_step_sft(): SFT-style training step """ job: Job = field(default_factory=Job) @@ -114,9 +122,17 @@ def __post_init__(self): os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True" logger.info("Compiling loss") self.loss = torch.compile(self.loss) + # ADD THIS LINE: Initialize ForgeEngine here (during __init__) to ensure + # init_distributed is called before Monarch initializes its process group + self._setup_engine() @endpoint async def setup(self): + # Already initialized in post_init__ + pass + + def _setup_engine(self): + """Initialize the ForgeEngine (non-endpoint helper for subclasses).""" # TODO: update ForgeEngine to not use ForgeJobConfig engine_config = {f.name: getattr(self, f.name) for f in fields(self)} for key in { @@ -127,17 +143,114 @@ async def setup(self): }: engine_config.pop(key) # Not part of job config self.engine = ForgeEngine(ForgeJobConfig(**engine_config)) + + # all ranks should record loss, except when PP=True. Then, only the last stage should record loss. + self.rank_should_record_loss = True + if ( + hasattr(self.engine, "pp_has_last_stage") + and not self.engine.pp_has_last_stage + ): + self.rank_should_record_loss = False + self.engine.checkpointer.load(step=self.step) self.engine.optimizers.zero_grad() + async def setup_metric_logger(self): + """Initialization happens in the main process. Here we just retrieve it.""" + mlogger = await get_or_create_metric_logger() + return mlogger + + def record_batch_metrics(self, data_metrics: list) -> None: + """Record metrics from batch. + + Since the dataloader creates new processes, we don't call `record_metric` in the dataset. + Instead, pop the metrics from the batch and record them here. + """ + for metric in data_metrics: + record_metric(metric.key, metric.value, metric.reduction) + def forward_backward( + self, + input_dict: dict[str, Tensor], + labels: Tensor, + skip_backward: bool = False, + ) -> Tensor: + """Forward and backward pass with PP and CP support. + + Args: + input_dict: Dict containing input tensors (e.g., {"tokens": tensor}) + labels: Target labels tensor + skip_backward: If True, skip backward pass (useful for eval) + + Returns: + Loss tensor + """ + model_parts = self.engine.model_parts + parallel_dims = self.engine.parallel_dims + + # apply context parallelism if cp is enabled + # ensure CP handles the separate freqs_cis buffer for each pp stage + inputs = input_dict["tokens"] + optional_context_parallel_ctx = ( + dist_utils.create_context_parallel_ctx( + cp_mesh=parallel_dims.world_mesh["cp"], + cp_buffers=[inputs, labels] + [m.freqs_cis for m in model_parts], + cp_seq_dims=[1, 1] + [0 for _ in model_parts], + cp_no_restore_buffers={inputs, labels}, + cp_rotate_method=self.engine.job_config.parallelism.context_parallel_rotate_method, + ) + if parallel_dims.cp_enabled + else None + ) + + if parallel_dims.pp_enabled: + # Pipeline Parallel forward / backward inside step() call + with self.engine.train_context(optional_context_parallel_ctx): + targets, losses = ( + (labels, []) if self.engine.pp_has_last_stage else (None, None) + ) + if self.engine.pp_has_first_stage: + self.engine.pp_schedule.step(inputs, target=targets, losses=losses) + else: + self.engine.pp_schedule.step(target=targets, losses=losses) + + # accumulate losses across pipeline microbatches + # TODO: PP+FSDP unexpectedly puts the loss back to the CPU + loss = ( + torch.sum(torch.stack(losses)).to(self.engine.device) + if self.engine.pp_has_last_stage + else torch.tensor(-1.0, device=self.engine.device) + ) + + # TODO: PP requires gradients enabled and cant deactive with no_grad + if skip_backward: + loss = loss.detach() + + else: + # Non-PP forward / backward + with self.engine.train_context(optional_context_parallel_ctx): + assert len(model_parts) == 1 + with self.engine.maybe_enable_amp: + pred = model_parts[0](inputs) + loss = self.engine.loss_fn(pred, labels) + # need to free to before bwd to avoid peaking memory + del pred + + # Only run backward if requested. Useful for eval. + if not skip_backward: + loss.backward() + + return loss + + def forward_backward_rl( self, inputs: dict[str, Tensor], targets: dict[str, Tensor] ) -> Tensor: + """Forward and backward for RL (uses custom self.loss instead of engine.loss_fn).""" model_parts = self.engine.model_parts parallel_dims = self.engine.parallel_dims optional_context_parallel_ctx = None if parallel_dims.pp_enabled: - raise NotImplementedError("PP not implemented yet") + raise NotImplementedError("PP not implemented yet for RL") else: with self.engine.train_context(optional_context_parallel_ctx): assert len(model_parts) == 1 @@ -148,10 +261,35 @@ def forward_backward( loss.backward() return loss + def train_step_sft(self, batch: dict[str, Tensor]) -> None: + """Execute a single SFT training step. + + This method is designed for supervised fine-tuning where the batch + contains "tokens" and "labels". It performs forward/backward pass, + and optimizer step. + + Args: + batch: Batch dict containing "tokens" and "labels" + """ + labels = batch.pop("labels") + loss = self.forward_backward(batch, labels) + + if self.rank_should_record_loss: + loss_val = loss.item() + record_metric("TitanTrainer/train_step_sft/loss", loss_val, Reduce.MEAN) + logger.info( + f"step {self.step} / {self.num_training_steps} | Loss: {loss_val}" + ) + + self.engine.optimizers.step() + self.engine.lr_schedulers.step() + self.step += 1 + @endpoint async def train_step( self, inputs: list[dict[str, Tensor]], targets: list[dict[str, Tensor]] ) -> float: + """Training step for RL workloads (endpoint version).""" t = Tracer("rl_trainer_perf/step", timer="gpu", track_memory=True) t.start() @@ -161,7 +299,7 @@ async def train_step( batch_to_device(local_inputs, self.engine.device) batch_to_device(local_targets, self.engine.device) - loss = self.forward_backward(local_inputs, local_targets) + loss = self.forward_backward_rl(local_inputs, local_targets) torch.distributed.all_reduce(loss) t.step("forward_backward") @@ -178,11 +316,6 @@ async def train_step( loss = loss.detach().item() record_metric("rl_trainer/loss", loss, Reduce.MEAN) - # These are placeholder values until the loss function exposes these metrics - # record_metric("rl_trainer/step/avg_kl_divergence", 0.0, Reduce.MEAN) - # record_metric("rl_trainer/step/std_kl_divergence", 0.0, Reduce.STD) - # record_metric("rl_trainer/step/avg_policy_entropy", 0.0, Reduce.MEAN) - self.step += 1 self.engine.checkpointer.save( curr_step=self.step, From 7371250b9d92e95cf8baecffd6c8ade89088fa1c Mon Sep 17 00:00:00 2001 From: Hossein Kavianihamedani Date: Fri, 26 Dec 2025 10:50:09 -0800 Subject: [PATCH 3/5] Move compile config from training to top-level - Move compile from training.compile to compile.enable in llama3_8b.yaml - Move compile from training.compile to compile.enable in qwen3_8b.yaml This aligns the YAML config structure with ForgeJobConfig's expected schema, where compile is a separate top-level dataclass, not nested under training. --- apps/sft/llama3_8b.yaml | 15 +++++++++------ apps/sft/qwen3_8b.yaml | 7 ++++--- 2 files changed, 13 insertions(+), 9 deletions(-) diff --git a/apps/sft/llama3_8b.yaml b/apps/sft/llama3_8b.yaml index 669aaf2df..98d997496 100644 --- a/apps/sft/llama3_8b.yaml +++ b/apps/sft/llama3_8b.yaml @@ -12,7 +12,7 @@ model_name: "meta-llama/Meta-Llama-3.1-8B-Instruct" model: name: llama3 flavor: 8B - hf_assets_path: /home/hosseinkh/models/Meta-Llama-3.1-8B-Instruct + hf_assets_path: hf://${model_name} processes: procs: 8 @@ -56,19 +56,22 @@ parallelism: checkpoint: enable: true - folder: /home/hosseinkh/models/Meta-Llama-3.1-8B-Instruct/checkpoint # The folder to save checkpoints to. - initial_load_path: /home/hosseinkh/models/Meta-Llama-3.1-8B-Instruct # The path to load the initial checkpoint from. Ignored if `folder` exists. + folder: ./checkpoint # The folder to save checkpoints to. + initial_load_path: hf://${model_name} # The path to load the initial checkpoint from. Ignored if `folder` exists. initial_load_in_hf: true # If true, interpret initial_load_path as a HuggingFace model repo last_save_in_hf: true - interval: 50 + interval: 500 async_mode: "disabled" activation_checkpoint: mode: selective selective_ac_option: op -metric_logging: {} - +metric_logging: + wandb: + project: sft-training + group: sft_exp_${oc.env:USER} + logging_mode: global_reduce # global_reduce, per_rank_reduce, per_rank_no_reduce # profiling: # enable_profiling: false diff --git a/apps/sft/qwen3_8b.yaml b/apps/sft/qwen3_8b.yaml index a8d2244a9..b8303b873 100644 --- a/apps/sft/qwen3_8b.yaml +++ b/apps/sft/qwen3_8b.yaml @@ -25,16 +25,17 @@ optimizer: lr_scheduler: warmup_steps: 200 +compile: + enable: false + training: local_batch_size: 8 seq_len: 2048 max_norm: 1.0 - steps: 1000 - compile: false + steps: 300 datasets: - path: "yahma/alpaca-cleaned" split: "train[:95%]" - eval: eval_every_n_steps: 50 # null = disabled max_eval_steps: null # null = run until epoch completes From f2497674a7b630a809ad665538efe7083a5488cc Mon Sep 17 00:00:00 2001 From: Hossein Kavianihamedani Date: Fri, 26 Dec 2025 10:56:31 -0800 Subject: [PATCH 4/5] Remove unnecessary files - Update yamls without compile under training --- apps/sft/OPTION2_INTEGRATION_GUIDE.md | 172 -------------- apps/sft/hf_checkpoint.py | 228 ------------------ apps/sft/llama3_8b.yaml | 2 +- apps/sft/qwen3_8b.yaml | 5 +- docs/TRAINER_REFACTORING.md | 322 -------------------------- 5 files changed, 4 insertions(+), 725 deletions(-) delete mode 100644 apps/sft/OPTION2_INTEGRATION_GUIDE.md delete mode 100644 apps/sft/hf_checkpoint.py delete mode 100644 docs/TRAINER_REFACTORING.md diff --git a/apps/sft/OPTION2_INTEGRATION_GUIDE.md b/apps/sft/OPTION2_INTEGRATION_GUIDE.md deleted file mode 100644 index b3e3e6781..000000000 --- a/apps/sft/OPTION2_INTEGRATION_GUIDE.md +++ /dev/null @@ -1,172 +0,0 @@ -# Example: How to integrate HuggingFace checkpoint saving into main.py - -## Option 2: Integration Steps - -### Step 1: Import the HF checkpoint module in main.py - -Add to the imports section (around line 13-30): -```python -from forge.apps.sft.hf_checkpoint import save_hf_checkpoint, load_hf_checkpoint -``` - -### Step 2: Replace the existing checkpoint save call - -Find this code in the `train()` method (around line 456): -```python -self.checkpointer.save( - curr_step=self.current_step, - last_step=self.current_step == self.num_training_steps, -) -``` - -Replace it with: -```python -# Save checkpoint using HuggingFace format (bypasses DCP KeyError bug) -checkpoint_config = self.job_config.checkpoint -if (checkpoint_config.enable and - self.current_step % checkpoint_config.interval == 0): - - save_hf_checkpoint( - model=self.model[0], # First model in list - tokenizer=self.tokenizer, - optimizer=self.optimizer.optimizers[0] if self.optimizer else None, - lr_scheduler=self.lr_scheduler.schedulers[0] if self.lr_scheduler else None, - checkpoint_dir=checkpoint_config.folder, - step=self.current_step, - rank=self._rank, - save_optimizer=True, - ) -``` - -### Step 3: Update the config to enable checkpointing - -In `llama3_8b.yaml`, change: -```yaml -checkpoint: - enable: true # ✅ Re-enable checkpointing! - folder: /home/hosseinkh/models/Meta-Llama-3.1-8B-Instruct/checkpoint - initial_load_path: /home/hosseinkh/models/Meta-Llama-3.1-8B-Instruct/ - initial_load_in_hf: true - last_save_in_hf: true # ✅ Now this will work! - interval: 500 - async_mode: "disabled" -``` - -### Step 4: Optional - Load from HF checkpoint at startup - -In the `setup()` method (around line 138), replace: -```python -self.checkpointer.load(step=self.current_step) -``` - -With: -```python -# Load from HuggingFace checkpoint if exists -checkpoint_config = self.job_config.checkpoint -if checkpoint_config.enable and Path(checkpoint_config.folder).exists(): - try: - loaded_step = load_hf_checkpoint( - model=self.model[0], - tokenizer=self.tokenizer, - optimizer=self.optimizer.optimizers[0] if self.optimizer else None, - lr_scheduler=self.lr_scheduler.schedulers[0] if self.lr_scheduler else None, - checkpoint_dir=checkpoint_config.folder, - step=None, # Load latest - rank=self._rank, - ) - self.current_step = loaded_step - logger.info(f"Resumed from step {loaded_step}") - except Exception as e: - logger.warning(f"Could not load checkpoint: {e}, starting from scratch") -``` - ---- - -## How It Works - -### The Problem with PyTorch DCP: -``` -HuggingFace Model (plain state_dict) - ↓ -FSDP Wrapping (changes parameter IDs) - ↓ -Training with new parameter IDs - ↓ -PyTorch DCP Save (tries to map back to original FQNs) - ❌ KeyError: 289 - parameter ID not in fqn_pid_mapping -``` - -### The HuggingFace Solution: -``` -FSDP Model (sharded across GPUs) - ↓ -Gather Full State Dict (FSDP.state_dict_type) - ↓ -Load into Unwrapped Model - ↓ -HuggingFace save_pretrained (safetensors format) - ✅ Works perfectly! -``` - ---- - -## Key Features of HF Checkpoint - -1. **FSDP-Compatible**: Uses `FSDP.state_dict_type()` to properly gather sharded weights -2. **Rank-0 Only**: Only rank 0 saves to avoid race conditions -3. **Barriers**: Proper synchronization across all ranks -4. **Unwrapping**: Recursively unwraps FSDP layers to get original model -5. **SafeTensors**: Uses modern safetensors format -6. **Optimizer State**: Separately saves optimizer/scheduler as PyTorch .pt files -7. **Metadata**: Saves training step and config for easy resumption - ---- - -## File Structure After Saving - -``` -checkpoint_dir/ -├── step_500/ -│ ├── model.safetensors # ✅ Model weights (HF format) -│ ├── config.json # ✅ Model config -│ ├── tokenizer.json # ✅ Tokenizer -│ ├── tokenizer_config.json # ✅ Tokenizer config -│ ├── optimizer.pt # ✅ Optimizer + LR scheduler -│ └── training_metadata.pt # ✅ Step number, etc. -├── step_1000/ -│ └── ... -``` - ---- - -## Benefits Over PyTorch DCP - -| Feature | PyTorch DCP | HuggingFace | -|---------|-------------|-------------| -| FSDP + HF Checkpoint | ❌ KeyError: 289 | ✅ Works | -| Format Compatibility | DCP-specific | ✅ Universal HF format | -| Easy Loading | Complex | ✅ `from_pretrained()` | -| Safetensors Support | Limited | ✅ Native | -| Resume Training | ✅ Yes | ✅ Yes | -| Size | Larger | Smaller (compressed) | - ---- - -## Testing - -After implementing, test with: -```bash -cd /home/hosseinkh/forge_updated/forge -python -m apps.sft.main --config apps/sft/llama3_8b.yaml -``` - -At step 500, you should see: -``` -INFO:forge.apps.sft.hf_checkpoint:Saving HuggingFace checkpoint to /home/hosseinkh/models/Meta-Llama-3.1-8B-Instruct/checkpoint/step_500 -INFO:forge.apps.sft.hf_checkpoint:✅ Model saved to /home/hosseinkh/models/Meta-Llama-3.1-8B-Instruct/checkpoint/step_500 -INFO:forge.apps.sft.hf_checkpoint:✅ Tokenizer saved to /home/hosseinkh/models/Meta-Llama-3.1-8B-Instruct/checkpoint/step_500 -INFO:forge.apps.sft.hf_checkpoint:✅ Optimizer state saved to /home/hosseinkh/models/Meta-Llama-3.1-8B-Instruct/checkpoint/step_500/optimizer.pt -INFO:forge.apps.sft.hf_checkpoint:🎉 Checkpoint saved successfully at step 500 -``` - -No more KeyError! 🎉 diff --git a/apps/sft/hf_checkpoint.py b/apps/sft/hf_checkpoint.py deleted file mode 100644 index 77799cf88..000000000 --- a/apps/sft/hf_checkpoint.py +++ /dev/null @@ -1,228 +0,0 @@ -# Copyright (c) Meta Platforms, Inc. and affiliates. -# All rights reserved. -# -# This source code is licensed under the BSD-style license found in the -# LICENSE file in the root directory of this source tree. - -""" -HuggingFace-native checkpoint saving that bypasses PyTorch DCP. - -This module provides checkpoint saving functionality that works with FSDP -by gathering the full model state and using HuggingFace's native save_pretrained. -""" - -import logging -import os -from pathlib import Path -from typing import Optional - -import torch -import torch.distributed as dist -from torch.distributed.fsdp import FullyShardedDataParallel as FSDP -from torch.distributed.fsdp import StateDictType, FullStateDictConfig - -logger = logging.getLogger(__name__) - - -def save_hf_checkpoint( - model: torch.nn.Module, - tokenizer, - optimizer: Optional[torch.optim.Optimizer], - lr_scheduler: Optional[torch.optim.lr_scheduler._LRScheduler], - checkpoint_dir: str, - step: int, - rank: int = 0, - save_optimizer: bool = True, -) -> None: - """ - Save checkpoint in HuggingFace format, bypassing PyTorch DCP. - - This function gathers the full model state from FSDP shards and saves - it using HuggingFace's native format, avoiding the KeyError bug in DCP. - - Args: - model: The model to save (can be FSDP-wrapped) - tokenizer: The tokenizer to save - optimizer: Optional optimizer to save - lr_scheduler: Optional LR scheduler to save - checkpoint_dir: Base directory for checkpoints - step: Current training step - rank: Process rank (only rank 0 saves) - save_optimizer: Whether to save optimizer state - """ - # Create checkpoint directory - ckpt_path = Path(checkpoint_dir) / f"step_{step}" - - if rank == 0: - ckpt_path.mkdir(parents=True, exist_ok=True) - logger.info(f"Saving HuggingFace checkpoint to {ckpt_path}") - - # Wait for directory creation - if dist.is_initialized(): - dist.barrier() - - # Handle FSDP model - if isinstance(model, FSDP): - # Configure FSDP to gather full state dict - save_policy = FullStateDictConfig(offload_to_cpu=True, rank0_only=True) - with FSDP.state_dict_type( - model, - StateDictType.FULL_STATE_DICT, - state_dict_config=save_policy, - ): - model_state = model.state_dict() - else: - model_state = model.state_dict() - - # Only rank 0 saves to avoid race conditions - if rank == 0: - # Get the unwrapped model for HF saving - if isinstance(model, FSDP): - # Access the original model inside FSDP wrapper - unwrapped_model = model._fsdp_wrapped_module - # If it's still wrapped (nested FSDP), keep unwrapping - while hasattr(unwrapped_model, '_fsdp_wrapped_module'): - unwrapped_model = unwrapped_model._fsdp_wrapped_module - else: - unwrapped_model = model - - # Load the full state dict into unwrapped model - unwrapped_model.load_state_dict(model_state, strict=True) - - # Save using HuggingFace's native method - try: - unwrapped_model.save_pretrained( - ckpt_path, - safe_serialization=True, # Use safetensors format - ) - logger.info(f"✅ Model saved to {ckpt_path}") - except Exception as e: - logger.error(f"❌ Failed to save model: {e}") - raise - - # Save tokenizer - if tokenizer is not None: - try: - tokenizer.save_pretrained(ckpt_path) - logger.info(f"✅ Tokenizer saved to {ckpt_path}") - except Exception as e: - logger.warning(f"⚠️ Failed to save tokenizer: {e}") - - # Save optimizer and scheduler (as PyTorch native - no HF format for these) - if save_optimizer and optimizer is not None: - try: - opt_path = ckpt_path / "optimizer.pt" - optimizer_state = { - 'optimizer': optimizer.state_dict(), - 'step': step, - } - if lr_scheduler is not None: - optimizer_state['lr_scheduler'] = lr_scheduler.state_dict() - - torch.save(optimizer_state, opt_path) - logger.info(f"✅ Optimizer state saved to {opt_path}") - except Exception as e: - logger.warning(f"⚠️ Failed to save optimizer: {e}") - - # Save training metadata - try: - meta_path = ckpt_path / "training_metadata.pt" - metadata = { - 'step': step, - 'model_config': unwrapped_model.config.to_dict() if hasattr(unwrapped_model, 'config') else {}, - } - torch.save(metadata, meta_path) - logger.info(f"✅ Metadata saved to {meta_path}") - except Exception as e: - logger.warning(f"⚠️ Failed to save metadata: {e}") - - logger.info(f"🎉 Checkpoint saved successfully at step {step}") - - # Synchronize all ranks - if dist.is_initialized(): - dist.barrier() - - -def load_hf_checkpoint( - model: torch.nn.Module, - tokenizer, - optimizer: Optional[torch.optim.Optimizer], - lr_scheduler: Optional[torch.optim.lr_scheduler._LRScheduler], - checkpoint_dir: str, - step: Optional[int] = None, - rank: int = 0, -) -> int: - """ - Load checkpoint from HuggingFace format. - - Args: - model: The model to load into - tokenizer: The tokenizer (not modified, HF loads from config) - optimizer: Optional optimizer to load state into - lr_scheduler: Optional LR scheduler to load state into - checkpoint_dir: Base directory for checkpoints - step: Specific step to load (None = latest) - rank: Process rank - - Returns: - The step number that was loaded - """ - ckpt_base = Path(checkpoint_dir) - - # Find checkpoint to load - if step is None: - # Find latest checkpoint - checkpoints = sorted([d for d in ckpt_base.glob("step_*") if d.is_dir()]) - if not checkpoints: - logger.warning(f"No checkpoints found in {checkpoint_dir}") - return 0 - ckpt_path = checkpoints[-1] - step = int(ckpt_path.name.split("_")[1]) - else: - ckpt_path = ckpt_base / f"step_{step}" - if not ckpt_path.exists(): - raise FileNotFoundError(f"Checkpoint not found: {ckpt_path}") - - if rank == 0: - logger.info(f"Loading HuggingFace checkpoint from {ckpt_path}") - - # Get unwrapped model - if isinstance(model, FSDP): - unwrapped_model = model._fsdp_wrapped_module - while hasattr(unwrapped_model, '_fsdp_wrapped_module'): - unwrapped_model = unwrapped_model._fsdp_wrapped_module - else: - unwrapped_model = model - - # Load model using HuggingFace's from_pretrained - try: - from transformers import AutoModelForCausalLM - loaded_model = AutoModelForCausalLM.from_pretrained( - ckpt_path, - torch_dtype=unwrapped_model.dtype if hasattr(unwrapped_model, 'dtype') else torch.float32, - ) - unwrapped_model.load_state_dict(loaded_model.state_dict(), strict=True) - if rank == 0: - logger.info(f"✅ Model loaded from {ckpt_path}") - except Exception as e: - logger.error(f"❌ Failed to load model: {e}") - raise - - # Load optimizer and scheduler - if optimizer is not None: - opt_path = ckpt_path / "optimizer.pt" - if opt_path.exists(): - try: - optimizer_state = torch.load(opt_path, map_location='cpu') - optimizer.load_state_dict(optimizer_state['optimizer']) - if lr_scheduler is not None and 'lr_scheduler' in optimizer_state: - lr_scheduler.load_state_dict(optimizer_state['lr_scheduler']) - if rank == 0: - logger.info(f"✅ Optimizer state loaded from {opt_path}") - except Exception as e: - logger.warning(f"⚠️ Failed to load optimizer: {e}") - - if rank == 0: - logger.info(f"🎉 Checkpoint loaded successfully from step {step}") - - return step diff --git a/apps/sft/llama3_8b.yaml b/apps/sft/llama3_8b.yaml index 98d997496..48b59e9b1 100644 --- a/apps/sft/llama3_8b.yaml +++ b/apps/sft/llama3_8b.yaml @@ -33,7 +33,7 @@ training: local_batch_size: 8 seq_len: 2048 max_norm: 1.0 - steps: 300 + steps: 1000 datasets: - path: "yahma/alpaca-cleaned" split: "train[:95%]" diff --git a/apps/sft/qwen3_8b.yaml b/apps/sft/qwen3_8b.yaml index b8303b873..62079854d 100644 --- a/apps/sft/qwen3_8b.yaml +++ b/apps/sft/qwen3_8b.yaml @@ -32,12 +32,13 @@ training: local_batch_size: 8 seq_len: 2048 max_norm: 1.0 - steps: 300 + steps: 1000 datasets: - path: "yahma/alpaca-cleaned" split: "train[:95%]" + eval: - eval_every_n_steps: 50 # null = disabled + eval_every_n_steps: 10 # null = disabled max_eval_steps: null # null = run until epoch completes datasets: - path: "yahma/alpaca-cleaned" diff --git a/docs/TRAINER_REFACTORING.md b/docs/TRAINER_REFACTORING.md deleted file mode 100644 index 3bdc02ec5..000000000 --- a/docs/TRAINER_REFACTORING.md +++ /dev/null @@ -1,322 +0,0 @@ -# Trainer Refactoring RFC: Unified TitanTrainer - -## Overview - -This RFC defines the target class hierarchy for unifying SFT and RL trainers. - ---- - -## Target Class Hierarchy - -``` -ForgeActor (base actor) - │ - └── TitanTrainer (generic trainer with ForgeEngine composition) - │ - └── ForgeSFTRecipe (SFT-specific: data loading, eval, training loop) - -ForgeEngine (standalone training engine - used via composition, not inheritance) -``` - ---- - -## TitanTrainer: The Base Trainer - -```python -@dataclass -class TitanTrainer(ForgeActor): - """Base trainer that wraps ForgeEngine with composition pattern. - - Provides reusable components for any training workload: - - _setup_engine(): Engine initialization - - rank_should_record_loss: Only last PP stage records loss - - record_batch_metrics(): Generic batch metric recording - - setup_metric_logger(): Metric logger initialization - - forward_backward(): Forward/backward with PP+CP support (uses engine.loss_fn) - - forward_backward_rl(): Forward/backward for RL (uses custom self.loss) - - train_step_sft(): SFT-style training step - - train_step(): RL-style training step (endpoint) - - push_weights(): Push weights to torchstore - - cleanup(): Cleanup resources - - Uses COMPOSITION: access ForgeEngine via self.engine.X - """ - - # === ForgeJobConfig fields === - job: Job = field(default_factory=Job) - model: Model = field(default_factory=Model) - optimizer: Optimizer = field(default_factory=Optimizer) - lr_scheduler: LRScheduler = field(default_factory=LRScheduler) - training: Training = field(default_factory=Training) - parallelism: Parallelism = field(default_factory=Parallelism) - checkpoint: Checkpoint = field(default_factory=Checkpoint) - activation_checkpoint: ActivationCheckpoint = field(default_factory=ActivationCheckpoint) - compile: Compile = field(default_factory=Compile) - quantize: Quantize = field(default_factory=Quantize) - comm: Comm = field(default_factory=Comm) - memory_estimation: MemoryEstimation = field(default_factory=MemoryEstimation) - - # === Non-JobConfig fields === - loss: Callable = lambda logits, **targets: logits # Custom loss for RL - - # === Engine (set during setup) === - engine: ForgeEngine # Composition - not inheritance! - rank_should_record_loss: bool - - # === Lifecycle Methods === - - @endpoint - async def setup(self): - self._setup_engine() - - def _setup_engine(self): - """Initialize ForgeEngine. Non-endpoint helper for subclasses.""" - self.engine = ForgeEngine(ForgeJobConfig(...)) - # Set rank_should_record_loss based on PP stage - self.rank_should_record_loss = True - if hasattr(self.engine, "pp_has_last_stage") and not self.engine.pp_has_last_stage: - self.rank_should_record_loss = False - - @endpoint - async def cleanup(self) -> None: - if self.engine.checkpointer: - self.engine.checkpointer.close() - - # === Reusable Utilities === - - async def setup_metric_logger(self): - """Get or create metric logger.""" - return await get_or_create_metric_logger() - - def record_batch_metrics(self, data_metrics: list): - """Record metrics from batch.""" - for metric in data_metrics: - record_metric(metric.key, metric.value, metric.reduction) - - # === Forward/Backward Passes === - - def forward_backward( - self, - input_dict: dict[str, Tensor], - labels: Tensor, - skip_backward: bool = False - ) -> Tensor: - """Forward/backward for SFT with PP+CP support. - - Uses engine.loss_fn (built-in cross-entropy loss). - Supports Pipeline Parallelism and Context Parallelism. - """ - # ... PP and CP handling ... - loss = self.engine.loss_fn(pred, labels) - if not skip_backward: - loss.backward() - return loss - - def forward_backward_rl( - self, - inputs: dict[str, Tensor], - targets: dict[str, Tensor] - ) -> Tensor: - """Forward/backward for RL with custom loss. - - Uses self.loss (custom RL loss function passed to trainer). - """ - logits = model(**inputs) - loss = self.loss(logits, **targets) # Custom loss! - loss.backward() - return loss - - # === Training Steps === - - def train_step_sft(self, batch: dict[str, Tensor]) -> None: - """SFT training step. Called from internal training loop. - - Extracts labels, calls forward_backward, logs metrics, optimizer step. - """ - labels = batch.pop("labels") - loss = self.forward_backward(batch, labels) - - if self.rank_should_record_loss: - record_metric("loss", loss.item(), Reduce.MEAN) - - self.engine.optimizers.step() - self.engine.lr_schedulers.step() - - @endpoint - async def train_step( - self, - inputs: list[dict[str, Tensor]], - targets: list[dict[str, Tensor]] - ) -> float: - """RL training step endpoint. Called from external control loop. - - For GRPO and similar algorithms where data comes from external rollouts. - """ - loss = self.forward_backward_rl(inputs[self.engine.dp_rank], targets[self.engine.dp_rank]) - torch.distributed.all_reduce(loss) - - self.engine.optimizers.step() - self.engine.optimizers.zero_grad() - self.engine.lr_schedulers.step() - - return loss.item() -``` - ---- - -## ForgeSFTRecipe: SFT-Specific Recipe - -```python -class ForgeSFTRecipe(TitanTrainer): - """SFT Recipe that adds data loading, eval, and training loop. - - INHERITS from TitanTrainer: - ✓ _setup_engine(), rank_should_record_loss - ✓ setup_metric_logger(), record_batch_metrics() - ✓ forward_backward(), train_step_sft() - - ADDS SFT-specific: - + __init__: Config handling (datasets, eval config) - + setup_data(): Create StatefulDataLoader with tokenizer - + evaluate(): Evaluation loop - + train(): Training loop with checkpointing and eval triggers - """ - - def __init__(self, config: DictConfig): - # Extract SFT-specific config (datasets, eval) - self.eval_config = config.get("eval", {}) - training_datasets = config.training.pop("datasets", None) - - # Initialize parent with ForgeJobConfig fields only - super().__init__(**titan_config) - - # Restore datasets for SFT use - self.job_config.training.datasets = training_datasets - - @endpoint - async def setup(self): - # Initialize engine (from TitanTrainer) - self._setup_engine() - - # Initialize metric logger (from TitanTrainer) - self.mlogger = await self.setup_metric_logger() - - # SFT-specific: Setup data loaders - self.train_dataloader = self.setup_data(self.job_config.training.datasets) - - # SFT-specific: Setup eval dataloaders - if self.validation_enabled: - for dataset_config in self.eval_config.datasets: - self.val_dataloaders[name] = self.setup_data([dataset_config]) - - def setup_data(self, dataset_configs: list[dict]) -> StatefulDataLoader: - """SFT-specific: Create dataloader with tokenizer.""" - tokenizer = HuggingFaceModelTokenizer(...) - dataset = sft_iterable_dataset(tokenizer, ...) - return StatefulDataLoader(dataset, ...) - - async def evaluate(self) -> None: - """SFT-specific: Evaluation loop.""" - for model_part in self.engine.model_parts: - model_part.eval() - - for batch in val_dataloader: - labels = batch.pop("labels") - # Use inherited forward_backward with skip_backward=True - loss = self.forward_backward(batch, labels, skip_backward=True) - - for model_part in self.engine.model_parts: - model_part.train() - - @endpoint - async def train(self) -> None: - """SFT-specific: Training loop with eval triggers.""" - self.engine.optimizers.zero_grad() - - while self.current_step < self.num_training_steps: - batch = next(self.train_dataloader) - - # Use inherited utilities - self.record_batch_metrics(batch.pop("metrics", [])) - - # Move to device - batch = {k: v.to(self.engine.device) for k, v in batch.items()} - - # Use inherited train_step_sft - self.train_step_sft(batch) - self.current_step += 1 - - # Eval periodically - if self.current_step % self.eval_every_n_steps == 0: - await self.evaluate() - - # Checkpoint - self.engine.checkpointer.save(curr_step=self.current_step, ...) -``` - ---- - -## GRPO Usage (External Control Loop) - -```python -# GRPO uses TitanTrainer directly (no subclass needed for now) -# Data/rollouts managed by external Orchestrator - -async def run_grpo(): - # Create trainer with custom GRPO loss - trainer = await TitanTrainer.options(...).as_actor( - loss=grpo_loss_fn, # Custom RL loss - model={"name": "llama3", ...}, - ) - await trainer.setup.call() - - # External control loop (Orchestrator) - for epoch in range(num_epochs): - # Generate rollouts (external) - rollouts = await generate_rollouts(...) - - # Train on rollouts using train_step endpoint - for batch in rollouts: - inputs = [{"input_ids": batch.ids, "attention_mask": batch.mask}] - targets = [{"advantages": batch.advs, "old_log_probs": batch.log_probs}] - loss = await trainer.train_step.call(inputs, targets) - - await trainer.cleanup.call() -``` - ---- - -## Key Design Decisions - -| Decision | Rationale | -|----------|-----------| -| **Composition for ForgeEngine** | TitanTrainer wraps engine via `self.engine.X` - avoids diamond inheritance | -| **Inheritance for Recipes** | ForgeSFTRecipe inherits from TitanTrainer - gets all utilities automatically | -| **Two forward_backward variants** | `forward_backward()` uses engine.loss_fn (SFT); `forward_backward_rl()` uses custom self.loss (RL) | -| **Endpoint vs Non-endpoint** | `train_step()` is endpoint for external control (RL); `train_step_sft()` is regular method for internal loop (SFT) | -| **`_setup_engine()` helper** | Non-endpoint method so subclasses can call it from their own setup endpoints | - ---- - -## What's Reusable in TitanTrainer - -| Component | SFT Uses | RL Uses | Notes | -|-----------|----------|---------|-------| -| `_setup_engine()` | ✓ | ✓ | Initialize ForgeEngine | -| `rank_should_record_loss` | ✓ | ✓ | PP stage handling | -| `setup_metric_logger()` | ✓ | ✓ | Metric logging | -| `record_batch_metrics()` | ✓ | - | Batch metrics from dataloader | -| `forward_backward()` | ✓ | - | Uses engine.loss_fn | -| `forward_backward_rl()` | - | ✓ | Uses custom self.loss | -| `train_step_sft()` | ✓ | - | Internal training step | -| `train_step()` endpoint | - | ✓ | External training step | -| `push_weights()` | - | ✓ | Push to torchstore | -| `cleanup()` | ✓ | ✓ | Resource cleanup | - ---- - -## Migration Path - -1. **PR 1 (Current)**: Add reusable components to TitanTrainer, have ForgeSFTRecipe inherit from TitanTrainer -2. **PR 2 (Future)**: Migrate GRPO to use TitanTrainer if beneficial -3. **PR 3 (Future)**: Add more reusable components (gradient clipping, etc.) From 0d525533bf295d37ed41a3af6f4deabab29db96e Mon Sep 17 00:00:00 2001 From: Hossein Kavianihamedani Date: Fri, 26 Dec 2025 10:58:53 -0800 Subject: [PATCH 5/5] Remove comment --- src/forge/controller/launcher.py | 1 - 1 file changed, 1 deletion(-) diff --git a/src/forge/controller/launcher.py b/src/forge/controller/launcher.py index 2a4a470d5..2f7e93aec 100644 --- a/src/forge/controller/launcher.py +++ b/src/forge/controller/launcher.py @@ -133,7 +133,6 @@ async def initialize(self) -> None: # HostMesh currently requires explicit configuration # of the underlying transport from client to mesh. # This can be removed in the future once this has been removed. - # Changed TCP this to TcpWithHostname for the slurm launcher configure(default_transport=ChannelTransport.TcpWithHostname) async def get_allocator(self, name: str, num_hosts: int) -> tuple[Any, Any, str]: