Merge branch 'main' into replica

Author: Allen Wang

apps/sft/llama3_8b.yaml

@@ -31,7 +31,19 @@ training:
   max_norm: 1.0
   steps: 1000
   compile: false
-  dataset: "c4"
+
+validation:
+  local_batch_size: 1
+  freq: -1  # Change to a positive number to enable validation
+  steps: 200  # Max steps to run validation. Validation disabled if negative.
+
+dataset:
+  path: yahma/alpaca-cleaned
+  split: train[:95%]
+
+dataset_val:
+  path: yahma/alpaca-cleaned
+  split: train[95%:]
 
 parallelism:
   data_parallel_replicate_degree: 1

apps/sft/main.py

@@ -18,9 +18,11 @@
 from forge.data.datasets.packed import PackedDataset, TextPacker
 from forge.data.datasets.sft_dataset import AlpacaToMessages, sft_iterable_dataset
 from forge.data.tokenizer import HuggingFaceModelTokenizer
+from forge.data.utils import batch_to_device, CROSS_ENTROPY_IGNORE_IDX
 
 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
@@ -30,6 +32,7 @@
 from torchtitan.experiments.forge.job_config import ForgeJobConfig
 from tqdm import tqdm
 
+
 # stubs for now
 Checkpointer = Any
 Dataloader = Any
@@ -63,7 +66,16 @@ def __init__(self, job_config: ForgeJobConfig):
         self.metric_logger = None  # TODO: fix this
 
     def setup(self):
-        self.train_dataloader = self.setup_data()
+        self.train_dataloader = self.setup_data(
+            self.job_config.dataset,
+            batch_size=self.job_config.training.local_batch_size,
+        )
+
+        self.val_dataloader = self.setup_data(
+            self.job_config.dataset_val,
+            batch_size=self.job_config.validation.local_batch_size,
+        )
+
         # self.train_dataloader = self.setup_data(
         #     self.train_config.train_dataset_config,
         #     self.train_config.train_dataloader_config,
@@ -79,7 +91,7 @@ def setup(self):
         # self.profiler = self.setup_profiler(self.train_config.profiler_config)
         # self.logger = self.setup_logger(self.train_config.logger_config)
 
-    def setup_data(self):
+    def setup_data(self, dataset_config, batch_size):
         tokenizer = HuggingFaceModelTokenizer(
             tokenizer_json_path=os.path.join(
                 self.job_config.model.hf_assets_path, "tokenizer.json"
@@ -95,8 +107,8 @@ def setup_data(self):
         dataset = sft_iterable_dataset(
             model_transform=tokenizer,
             message_transform=AlpacaToMessages(),
-            path="yahma/alpaca-cleaned",
-            split="train",
+            path=dataset_config.path,
+            split=dataset_config.split,
         )
         packer = TextPacker(padding_idx=0)
         dataset = PackedDataset(
@@ -106,7 +118,7 @@ def setup_data(self):
         )
         dataloader = StatefulDataLoader(
             dataset=dataset,
-            batch_size=self.job_config.training.local_batch_size,
+            batch_size=batch_size,
             collate_fn=partial(
                 collate_packed, mask_fn=packer.create_block_mask, device=self.device
             ),
@@ -119,7 +131,10 @@ def setup_data(self):
         return dataloader
 
     def forward_backward(
-        self, input_dict: dict[str, torch.Tensor], labels: torch.Tensor
+        self,
+        input_dict: dict[str, torch.Tensor],
+        labels: torch.Tensor,
+        do_backward: bool = True,
     ) -> torch.Tensor:
         model_parts = self.model_parts
         parallel_dims = self.parallel_dims
@@ -145,14 +160,16 @@ def forward_backward(
                 targets, losses = (
                     (labels, []) if self.pp_has_last_stage else (None, None)
                 )
+                if do_backward:
+                    pp_schedule_fn = self.pp_schedule.step
+                else:
+                    pp_schedule_fn = self.pp_schedule.eval
                 if self.pp_has_first_stage:
-                    self.pp_schedule.step(
+                    pp_schedule_fn(
                         inputs, target=targets, losses=losses, input_batch=inputs
                     )
                 else:
-                    self.pp_schedule.step(
-                        target=targets, losses=losses, input_batch=inputs
-                    )
+                    pp_schedule_fn(target=targets, losses=losses, input_batch=inputs)
 
             # accumulate losses across pipeline microbatches
             # TODO: PP+FSDP unexpectedly puts the loss back to the CPU
@@ -170,7 +187,8 @@ def forward_backward(
                     loss = self.loss_fn(pred, labels)
                 # need to free to before bwd to avoid peaking memory
                 del pred
-                loss.backward()
+                if do_backward:
+                    loss.backward()
 
         return loss
 
@@ -214,6 +232,52 @@ def train(self) -> None:
                 last_step=self.current_step == self.num_training_steps,
             )
 
+            if (
+                self.job_config.validation.freq > 0
+                and self.job_config.validation.steps > 0
+                and self.current_step % self.job_config.validation.freq == 0
+            ):
+                self.validate(self.job_config.validation.steps)
+
+    def validate(self, max_steps: int) -> None:
+        for m in self.model_parts:
+            m.eval()
+        total_val_loss = torch.tensor(0.0, device=self.device)
+        total_val_tokens = torch.tensor(0.0, device=self.device)
+        with torch.no_grad():
+            val_pbar = tqdm(self.val_dataloader, desc="Validation", leave=False)
+            for batch_idx, batch in enumerate(val_pbar):
+                if batch_idx >= max_steps:
+                    break
+                batch_to_device(batch, self.device)
+                current_num_tokens = (batch["labels"] != CROSS_ENTROPY_IGNORE_IDX).sum()
+                # Compute loss
+                labels = batch.pop("labels")
+                loss = self.forward_backward(batch, labels, do_backward=False)
+                val_loss = loss * current_num_tokens
+                total_val_loss += val_loss
+                total_val_tokens += current_num_tokens
+                # Update progress bar description with current average loss
+                avg_loss_so_far = (
+                    (total_val_loss / total_val_tokens).item()
+                    if total_val_tokens > 0
+                    else float("inf")
+                )
+                val_pbar.set_description(
+                    f"Running validation Loss: {avg_loss_so_far:.4f}"
+                )
+        # Aggregate validation metrics across all ranks
+        torch.distributed.all_reduce(total_val_loss)
+        torch.distributed.all_reduce(total_val_tokens)
+        avg_val_loss = (
+            (total_val_loss / total_val_tokens).item()
+            if total_val_tokens > 0
+            else float("inf")
+        )
+        for m in self.model_parts:
+            m.train()
+        print(f"\nValidation loss: {avg_val_loss}")
+
     def cleanup(self) -> None:
         if self.checkpointer:
             self.checkpointer.close()

src/forge/actors/policy.py

@@ -15,11 +15,13 @@
 import torch
 
 from monarch.actor import Actor, current_rank, endpoint, proc_mesh
+from torchstore import MultiProcessStore
+
+from torchstore._state_dict_utils import DELIM
 
 from vllm.engine.arg_utils import EngineArgs
 from vllm.entrypoints.utils import _validate_truncation_size
 from vllm.executor.multiproc_worker_utils import set_multiprocessing_worker_envs
-from vllm.inputs import TextPrompt, TokensPrompt
 from vllm.lora.request import LoRARequest
 from vllm.outputs import CompletionOutput
 from vllm.sampling_params import GuidedDecodingParams, RequestOutputKind, SamplingParams
@@ -37,6 +39,8 @@
 from vllm.v1.structured_output import StructuredOutputManager
 from vllm.worker.worker_base import WorkerWrapperBase
 
+from forge.data.sharding import VLLMSharding
+
 logger = logging.getLogger(__name__)
 
 
@@ -194,6 +198,7 @@ class Policy(Actor):
     enforce_eager: bool = False
     vllm_args: EngineArgs = None
     resources: int = 1
+    state_dict_key: str = "model_state_dict"
 
     def __post_init__(self):
         """Build vLLM Arguments
@@ -238,7 +243,8 @@ def __post_init__(self):
         assert self.vllm_args.parallel_config.world_size == self.resources
 
     @endpoint
-    async def setup(self):
+    async def setup(self, store: MultiProcessStore = None):
+        self.torchstore = store
         # TODO: remove ["gpus"] when monarch implements a flat rank
         self.rank = current_rank()["gpus"]
         self.worker = self.setup_worker()
@@ -247,10 +253,50 @@ async def setup(self):
     async def execute_model(self, schedule: SchedulerOutput):
         return self.worker.execute_model(schedule)
 
+    async def _load_tensor_parallel_state_dict(self, current_state_dict: dict):
+        """
+        Load full state dict from torchstore into tensor parallel model with deterministic sharding.
+        """
+
+        updated_count = 0
+        # setting explictly to llama3 for now as its our only use case
+        sharding = VLLMSharding(self.tensor_parallel_size, self.rank)
+
+        for param_name in current_state_dict.keys():
+            current_tensor = current_state_dict[param_name]
+
+            # Load the full tensor from torchstore
+            # TODO: only get the part of the tensor that is needed
+            stored_tensor = await self.torchstore.get(
+                f"{self.state_dict_key}{DELIM}{param_name}"
+            )
+            sharding.load_from_source_to_target(
+                param_name,
+                stored_tensor,
+                current_tensor,
+            )
+
+            updated_count += 1
+
     @endpoint
     async def update(self):
-        # TODO: add TorchStore support
-        pass
+        """Update model weights by reading state dict from torchstore"""
+
+        if self.torchstore is None:
+            raise Exception("No torchstore configured, skipping model update")
+
+        logger.debug(
+            f"Starting model update from torchstore with key: {self.state_dict_key}"
+        )
+
+        model = self.worker.model_runner.model
+        current_state_dict = model.state_dict()
+
+        logger.debug(f"Current state dict has {len(current_state_dict)} parameters")
+
+        await self._load_tensor_parallel_state_dict(current_state_dict)
+
+        logger.debug("Successfully updated model weights from torchstore")
 
     @endpoint
     async def setup_kv_cache(self):
@@ -286,6 +332,17 @@ async def setup_kv_cache(self):
     async def get_vllm_args(self):
         return self.vllm_args
 
+    @endpoint
+    async def get_model_params(self):
+        model = self.worker.model_runner.model
+        state_dict = {}
+
+        for name, param in model.named_parameters():
+            if "layers.0" not in name:
+                continue
+            state_dict[name] = param.cpu().detach()
+        return state_dict
+
     def setup_worker(self):
         """Build and Instantiate vLLM worker"""
         parallel_config = self.vllm_args.parallel_config