Data stream 2

.gitignore

@@ -4,6 +4,7 @@ wandb/*
 datasets/*
 output/*
 outputs/*
+data/*
 
 # Byte-compiled / optimized / DLL files
 __pycache__/

README.md

@@ -58,8 +58,20 @@ inference
 uv run python src/zeroband/inference.py @ configs/inference/debug.toml
 ```
 
+## Larger run
+
+For now you need to generate fake rollout data for testing. 
+
+```bash
+uv run python generate_fake_rollout.py @ configs/training/150M/A40.toml --data.path data/fake_rollout --optim.total_steps 1000
+```
+
+and then do the training on it
+
+```bash
+uv run torchrun --nproc_per_node=2 src/zeroband/train.py @ configs/training/150M/A40.toml --data.path data/fake_rollout 
+```
 
-...
 
 ## Checkpoints management
 

configs/inference/debug.toml

@@ -1,4 +1,5 @@
 name_model = "debugmodel"
 max_samples = 32
 batch_size = 8
+sample_per_file = 32
 dataset = "justus27/test-vcu"

configs/training/150M/H100-fast.toml

@@ -18,4 +18,4 @@ weight_decay = 0.24530252977858977
 
 [data]
 seq_length = 1024
-dataset_name_or_paths = "datasets/fineweb-edu"
+path = "datasets/fineweb-edu"

configs/training/debug.toml

@@ -11,4 +11,4 @@ total_steps = 4
 
 [data]
 fake = true
-
+timeout = 1

generate_fake_rollout.py

@@ -0,0 +1,29 @@
+# this is a script to generate fake rollout data for testing
+# should be removed once we have inference working
+# uv run python generate_fake_rollout.py  @ configs/training/150M/A40.toml --data.path data/fake_rollout --optim.total_steps 1000
+
+import os
+from pathlib import Path
+from pydantic_config import parse_argv
+from zeroband.train import Config
+
+from tests.conftest import _create_fake_rollout_parquet_file
+
+
+def main(config: Config):
+    path = Path(config.data.path)
+    os.makedirs(path, exist_ok=True)
+
+    num_files = 4
+    batch_size = config.optim.batch_size // num_files
+
+    _create_fake_rollout_parquet_file(
+        path, list(range(config.optim.total_steps)), num_files=num_files, batch_size=batch_size, seq_len=config.data.seq_length
+    )
+
+    print(f"Created test data at: {path}")
+
+
+if __name__ == "__main__":
+    config = Config(**parse_argv())  # type: ignore
+    main(config)

pyproject.toml

@@ -17,7 +17,8 @@ dependencies = [
     "pyarrow",
     "wandb",
     "vllm",
-    "jaxtyping"
+    "jaxtyping",
+    "beartype"
 ]
 
 

src/zeroband/inference.py

@@ -1,16 +1,35 @@
+import os
+import uuid
+from pydantic import model_validator
 from vllm import LLM, SamplingParams
 from pydantic_config import BaseConfig, parse_argv
+import vllm
 
 from zeroband.models import ModelName, name_to_hf_model
 
 from datasets import load_dataset
+import pyarrow as pa
+import pyarrow.parquet as pq
 
 
 class Config(BaseConfig):
     name_model: ModelName = "150M"
     dataset: str = "justus27/test-vcu"
     batch_size: int = 32
+    sample_per_file: int = 1024
     max_samples: int | None = None
+    output_path: str = "outputs"
+
+    @model_validator(mode="after")
+    def validate_bs_and_sample_per_file(self):
+        if self.sample_per_file % self.batch_size != 0:
+            raise ValueError("sample_per_file must be divisible by batch_size")
+        if self.max_samples is not None:
+            if self.max_samples % self.batch_size != 0:
+                raise ValueError("max_samples must be divisible by batch_size")
+            if self.max_samples < self.sample_per_file:
+                raise ValueError("max_samples must be greater than sample_per_file")
+        return self
 
 
 def fake_chat_template(messages):
@@ -28,6 +47,58 @@ def fake_chat_template(messages):
     return formatted_prompts
 
 
+pa_schema = pa.schema(
+    [
+        ("input_tokens", pa.list_(pa.int32())),
+        ("output_tokens", pa.list_(pa.int32())),
+        ("advantages", pa.float32()),
+        ("proofs", pa.binary()),
+        ("step", pa.int32()),
+    ]
+)
+
+
+def get_parquet_table(generated_tokens: list[vllm.RequestOutput], step: int) -> pa.Table:
+    # Initialize lists for each column
+    input_tokens_list = []
+    output_tokens_list = []
+    advantages_list = []
+    proofs_list = []
+    steps_list = []
+
+    # Process each RequestOutput
+    for request in generated_tokens:
+        # For each output in the request (handling top-n outputs)
+        for output in request.outputs:
+            # Input tokens are the prompt tokens
+            input_tokens_list.append(request.prompt_token_ids)
+
+            # Output tokens from the completion
+            output_tokens_list.append(output.token_ids)
+
+            # Initialize with 0 advantage as it's not part of RequestOutput
+            # You might want to modify this based on your advantage calculation
+            advantages_list.append(0)
+
+            # TODO: Add toploc proof
+            proofs_list.append("I am toploc proof, handcrafted by jack".encode())
+
+            # Add step
+            steps_list.append(step)
+
+    # Create PyArrow arrays
+    arrays = [
+        pa.array(input_tokens_list, type=pa.list_(pa.int32())),
+        pa.array(output_tokens_list, type=pa.list_(pa.int32())),
+        pa.array(advantages_list, type=pa.float32()),
+        pa.array(proofs_list, type=pa.binary()),
+        pa.array(steps_list, type=pa.int32()),
+    ]
+
+    # Create and return table
+    return pa.Table.from_arrays(arrays, schema=pa_schema)
+
+
 def main(config: Config):  # -> list[dict[str, Any]]:
     prompts = ["Write me a novel" for _ in range(5)]
 
@@ -41,6 +112,8 @@ def main(config: Config):  # -> list[dict[str, Any]]:
 
     max_samples = config.max_samples or len(dataset)
 
+    step = 0  # step will change once we have the update model api
+
     # Process batches
     for i in range(0, min(len(dataset), max_samples), config.batch_size):
         # Get batch
@@ -52,7 +125,14 @@ def main(config: Config):  # -> list[dict[str, Any]]:
         # Get tokenized inputs
         prompts = fake_chat_template(messages)
 
-        llm.generate(prompts, sampling_params)
+        generated_tokens = llm.generate(prompts, sampling_params)
+
+        table = get_parquet_table(generated_tokens, step)
+
+        step_path = f"{config.output_path}/step_{step}"
+        os.makedirs(step_path, exist_ok=True)
+
+        pq.write_table(table, f"{step_path}/{uuid.uuid4()}.parquet")
 
 
 if __name__ == "__main__":

src/zeroband/train.py

@@ -2,6 +2,7 @@
 import time
 from typing import TYPE_CHECKING, Literal
 
+from pydantic import model_validator
 import torch
 import torch.distributed as dist
 from torch.distributed.fsdp import fully_shard, MixedPrecisionPolicy  # type: ignore
@@ -64,12 +65,26 @@ class Config(BaseConfig):
     optim: OptimConfig = OptimConfig()
     train: TrainConfig
 
+    @model_validator(mode="after")
+    def check_batch_size(self):
+        if self.data.batch_size is None:
+            self.data.batch_size = self.optim.batch_size
+        assert self.optim.batch_size == self.data.batch_size, (
+            "The batch size in the config must be the same as the batch size in the data config."
+        )
+        return self
 
-def get_gradient_accumulation_steps(batch_size: int, micro_bs: int) -> int:
+
+def get_gradient_accumulation_steps(batch_size: int, micro_bs: int, data_workers: int) -> int:
     assert batch_size % world_info.local_world_size == 0
     batch_size = batch_size // world_info.local_world_size
 
     assert batch_size % micro_bs == 0, f"The micro batch size ({micro_bs}) must divide the number of samples on each GPU ({batch_size})."
+
+    assert batch_size % (data_workers * world_info.local_world_size) == 0, (
+        f"The batch size ({batch_size}) must be divisible by the number of data workers ({data_workers}) times the number of GPUs ({world_info.local_world_size})."
+    )
+
     return batch_size // micro_bs
 
 
@@ -88,7 +103,7 @@ def apply_fsdp(model: torch.nn.Module, reshard_after_forward: bool):
 def train(config: Config):
     # batch_size is the total batch size for all GPUs
 
-    gradient_accumulation_steps = get_gradient_accumulation_steps(config.optim.batch_size, config.train.micro_bs)
+    gradient_accumulation_steps = get_gradient_accumulation_steps(config.optim.batch_size, config.train.micro_bs, config.data.num_workers)
 
     model, tokenizer = get_model_and_tokenizer(config.name_model)
 
@@ -130,12 +145,13 @@ def train(config: Config):
             batch = next(train_dataloader_iterator)
 
             input_ids: Int[torch.Tensor, "batch seq"] = batch["input_ids"].to("cuda")
-            advantages: Float[torch.Tensor, "batch"] = batch["advantages"].to("cuda")
-            ref_logprobs: Float[torch.Tensor, "batch seq"] = batch["ref_logprobs"].to("cuda")
+            advantages: Float[torch.Tensor, "batch seq"] = batch["advantages"].to("cuda")
 
-            policy_logprobs = model(input_ids=input_ids).logits.contiguous()
+            policy_logprobs: Float[torch.Tensor, "batch seq vocab"] = model(input_ids=input_ids).logits.contiguous()
+            ref_logprobs: Float[torch.Tensor, "batch seq vocab"] = torch.ones_like(policy_logprobs)
 
             loss = grpo_loss(policy_logprobs, ref_logprobs, advantages) / gradient_accumulation_steps
+            # loss = policy_logprobs.sum() / gradient_accumulation_steps
 
             loss.backward()
             loss_batch += loss.detach().clone()