improve model debugging workflow

architectures/centralized/server/src/app.rs

@@ -7,18 +7,17 @@ use psyche_coordinator::{
     SOLANA_MAX_NUM_CLIENTS, TickResult,
 };
 
-use psyche_core::{FixedVec, NodeIdentity, Shuffle, SizedIterator, TokenSize};
+use psyche_core::{FixedVec, NodeIdentity, Shuffle, SizedIterator};
 use psyche_data_provider::{
-    DataProviderTcpServer, DataServerTui, LocalDataProvider, download_model_from_gcs_async,
-    download_model_repo_async,
+    DataProviderTcpServer, DataServerConfig, DataServerTui, LocalDataProvider,
+    download_model_from_gcs_async, download_model_repo_async,
 };
 use psyche_network::{ClientNotification, PublicKey, TcpServer};
 use psyche_tui::{
     CustomWidget, MaybeTui, TabbedWidget, logging::LoggerWidget, maybe_start_render_loop,
 };
 use psyche_watcher::{CoordinatorTui, OpportunisticData};
 use rand::RngCore;
-use serde::{Deserialize, Serialize};
 use std::collections::HashSet;
 use std::net::{Ipv4Addr, SocketAddr};
 use std::ops::ControlFlow;
@@ -147,20 +146,12 @@ impl App {
     }
 }
 
-#[derive(Serialize, Deserialize, Debug)]
-pub struct DataServerInfo {
-    pub dir: PathBuf,
-    pub token_size: TokenSize,
-    pub seq_len: usize,
-    pub shuffle_seed: [u8; 32],
-}
-
 impl App {
     #[allow(clippy::too_many_arguments)]
     pub async fn new(
         tui: bool,
         mut coordinator: Coordinator,
-        data_server_config: Option<DataServerInfo>,
+        data_server_config: Option<DataServerConfig>,
         coordinator_server_port: Option<u16>,
         save_state_dir: Option<PathBuf>,
         init_warmup_time: Option<u64>,
@@ -216,7 +207,7 @@ impl App {
                             anyhow!("Failed to parse training data server URL {:?}: {}", url, e)
                         })?;
                         let data_server_port = server_addr.port();
-                        let DataServerInfo {
+                        let DataServerConfig {
                             dir,
                             seq_len,
                             shuffle_seed,

architectures/centralized/server/src/main.rs

@@ -2,9 +2,10 @@ mod app;
 mod dashboard;
 
 use anyhow::{Context, Result};
-use app::{App, DataServerInfo};
+use app::App;
 use clap::{ArgAction, Parser};
 use psyche_coordinator::Coordinator;
+use psyche_data_provider::DataServerConfig;
 use psyche_tui::{
     LogOutput, ServiceInfo,
     logging::{MetricsDestination, OpenTelemetry, RemoteLogsDestination, TraceDestination},
@@ -111,7 +112,7 @@ struct RunArgs {
 fn load_config_state(
     state_path: PathBuf,
     data_config_path: Option<PathBuf>,
-) -> Result<(Coordinator, Option<DataServerInfo>)> {
+) -> Result<(Coordinator, Option<DataServerConfig>)> {
     let coordinator: Coordinator = toml::from_str(std::str::from_utf8(
         &std::fs::read(&state_path).with_context(|| {
             format!("failed to read coordinator state toml file {state_path:?}")
@@ -120,7 +121,7 @@ fn load_config_state(
 
     let data_server_config = match data_config_path {
         Some(config_path) => {
-            let mut data_config: DataServerInfo = toml::from_str(std::str::from_utf8(
+            let mut data_config: DataServerConfig = toml::from_str(std::str::from_utf8(
                 &std::fs::read(&config_path).with_context(|| {
                     format!("failed to read data server config toml file {config_path:?}")
                 })?,

nix/devShell.nix

@@ -69,6 +69,7 @@
                 cargo-watch
                 cargo-expand
                 cargo-nextest
+                cargo-all-features
 
                 self'.packages.solana-toolbox-cli
 

psyche-book/src/SUMMARY.md

@@ -23,6 +23,7 @@
   - [Running on-chain](./development/running-onchain.md)
   - [Running off-chain](./development/running-offchain.md)
   - [Implementing Models](./development/models.md)
+  - [make-tiny-init.py](./development/make-tiny-init.md)
   - [Python](./development/python.md)
   - [Secrets](./development/agenix.md)
   - [Building these docs](./development/book.md)

psyche-book/src/development/make-tiny-init.md

@@ -0,0 +1,25 @@
+## make-tiny-init reference
+
+The `make-tiny-init.py` script creates tiny versions of any psyche/hf/torchtitan-compatible model for fast iteration. It produces a tiny init checkpoint, a `config.json`, and a ready-to-use [run config](../enduser/run-config.md).
+
+### Model Input
+
+The script accepts one of the following model formats:
+
+```bash
+# HF repo slug: easiest way, downloads config + tokenizer
+nix develop .#python --command python scripts/make-tiny-init.py --repo Qwen/Qwen3-30B-A3B --save /tmp/test-qwen
+
+# from a raw HF config.json, use --tokenizer to include a tokenizer
+nix develop .#python --command python scripts/make-tiny-init.py  --config /path/to/config.json --save /tmp/tiny --tokenizer Qwen/Qwen3-30B-A3B
+
+# from a torchtitan job toml
+nix develop .#python --command python scripts/make-tiny-init.py  --toml /path/to/train.toml --save /tmp/tiny-run
+
+# from a psyche run config (reads model.LLM.checkpoint.Hub.repo_id)
+nix develop .#python --command python scripts/make-tiny-init.py  --psyche /path/to/state.toml --save /tmp/tiny-llama
+```
+
+Read the comments inside the make-tiny-init script for more information.
+
+> **Tip:** use `--preserve vocab_size` if you want to train on real tokenized data — otherwise the tiny vocab won't match your tokenizer.

psyche-book/src/development/models.md

@@ -4,40 +4,112 @@ This codebase includes a set of sample programs that let you design, implement,
 
 We currently only implement Llama and Deepseek (see `shared/modeling/src/models/`), but PRs are very welcome to add more architectures and model types.
 
-The `train` example, documented below, is useful to test how your model trains using AdamW vs DisTrO.
+The `train` binary, documented below, is useful to test how your model trains using AdamW vs DisTrO.
 
-## Running
+## Quickstart: test a new architecture
+
+The fastest way to test a new model architecture end-to-end:
+
+### 1. Create a tiny model
+
+Use `make-tiny-init.py` to shrink any model to minimal dimensions. This creates a minimally sized (probably a few million parameters) version that loads & trains in seconds, so you can iterate quickly on architecture code.
+
+You can pass an HF repo slug directly to have the script download the model config & tokenizer for you.
+
+```bash
+nix develop .#python --command python scripts/make-tiny-init.py \
+    --repo Qwen/Qwen3-30B-A3B \
+    --save ~/test-model
+```
+
+This will output a folder `~/test-model/checkpoint/`, with weights, config, and the tokenizer, and a file `~/test-model/train.toml`, a Psyche [run config](../enduser/run-config.md) with some default settings.
+
+This script shrinks all dimensions of the model to the smallest possible valid values, while keeping the model able to be trained with tensor parallelism = 8. More information is available at [make-tiny-init.py reference](./make-tiny-init.md).
+
+### 2. Train
+
+The generated run config uses dummy data and sensible defaults to test.
+
+```bash
+cd ~/test-model
+nix run .#train -- config ./train.toml
+```
+
+You can also pass runtime options:
 
 ```bash
-cargo run --example train -- ---help
+nix run .#train -- config ./train.toml \
+    --micro-batch 4 \
+    --device cuda:0
 ```
 
-You'll need a pre-tokenized dataset downloaded to your disk for training.
+> Since we're using `Dummy` data , we generate random tokens, so the loss won't go down, but we still run through the full forward/backward pass. Switch to a real dataset when you need to verify that training does make loss go down.
 
-> A PR is welcome to add an option to the trainer to use the HTTP data provider! You can refer to the http example in the data-provider crate for a sample implementation.
+You can edit `train.toml` directly or write your own from scratch. See the [run configuration](../enduser/run-config.md) section for more info.
 
-For a Llama 2 model, a pre-tokenized dataset to test with is available at [https://huggingface.co/datasets/emozilla/fineweb-10bt-tokenized-datatrove-llama2/](https://huggingface.co/datasets/emozilla/fineweb-10bt-tokenized-datatrove-llama2/tree/main).
-Psyche only needs the `.ds` files, and will load any/all `.ds` files in the specified folder - you can download just one for smaller tests.
+## Running without a config file
 
-If you've downloaded part or all of the above dataset into a folder `data/fineweb-10bt` inside the Psyche repo, you can start a simple training run on a 20m parameter Llama 2 model:
+For quick tests, you can skip the TOML file entirely and pass everything as command line args.
 
 ```bash
-cargo run --example train -- \
-    --model emozilla/llama2-20m-init \
+nix run .#train -- \
+    --model ./test-model/ \
     --data-path ./data/fineweb-10bt/ \
     --total-batch 2 \
-    --micro-batch 1
+    --micro-batch 1 \
+    --architecture HfLlama
 ```
 
+## Dumping the config from a live run
+
+If there's a run on-chain and you want to reproduce its configuration locally (e.g. to create a tiny model for debugging, or to start a new run with the same settings), use `dump-config`:
+
+```bash
+nix run .#run-manager dump-config \
+    --rpc https://api.devnet.solana.com \
+    --run-id <your_run_id>
+```
+
+This will print the full run config to stdout. You can redirect it to a file and use it directly, or run it through make-tiny-init.py to make a smaller version for developing with.
+
+```bash
+# dump the run's config
+nix run .#run-manager dump-config --rpc https://api.devnet.solana.com --run-id my-run > live-config.toml
+
+# create a tiny version of that run's model for local testing
+nix develop .#python --command python scripts/make-tiny-init.py \
+    --psyche live-config.toml \
+    --save /tmp/tiny-my-run
+```
+
+This is useful when you want to iterate on an architecture locally using the exact same model architecture and hypers as a production run.
+
+## Datasets
+
+The `train` binary accepts any psyche data source when using a run config..
+
+| Data location  | Description                                                       |
+| -------------- | ----------------------------------------------------------------- |
+| `Dummy`        | Random tokens. (loss will not go down)                            |
+| `Local`        | A directory of `.ds` files on disk                                |
+| `Http`         | Data served over HTTP (single URL, numbered files, or GCS bucket) |
+| `Preprocessed` | A HuggingFace dataset repo                                        |
+
+When running without a config, only local datasets are supported.
+
+For a Llama 2 model, a pre-tokenized dataset is available at [emozilla/fineweb-10bt-tokenized-datatrove-llama2](https://huggingface.co/datasets/emozilla/fineweb-10bt-tokenized-datatrove-llama2/tree/main).
+Psyche only needs the `.ds` files, and will load any/all `.ds` files in the specified folder — you can download just one for smaller tests.
+
 ## Adding a new model type
 
-The `train` example currently asssumes your model is a Llama or Deepseek v2/v3 model, and instantiates it via `(LlamaForCausalLM|DeepseekForCausalLM)::from_pretrained`.
+Both the main Psyche client and the `train` binary support a handful of model types:
+Llama 2/3 models, a Deepseek v2/v3 models, any HuggingFace Transformers-compatible model, or any TorchTitan-compatible model.
 
-We currently only support causal language models - to implement a new one, you can create a file similar to `llama_for_causal_lm` and implement your model, ensuring you provide a trait impl for `CausalLM`.
+They're instantiated via `(LlamaForCausalLM|DeepseekForCausalLM)::from_pretrained` or `(PythonCausalLM::new|PythonDistributedCausalLM::new)`.
 
-There's alpha-level support for models written in Python. See the [Python](./python.md) docs for more information.
+We currently only support causal language models — to implement a new one, you can create a file similar to `llama_for_causal_lm` and implement your model, ensuring you provide a trait impl for `CausalLM` - or, preferrably, add your model to [our TorchTitan fork](https://github.com/nousResearch/torchtitan). See the [Python](./python.md) docs for more information.
 
-You might also need to modify the data provider, if your data is structured in some way.
+You might also need to modify the data provider if the data your model requires is structured in some way.
 Since you're implementing the forward pass yourself, you can serve and interpret data passed from the data provider however you need.
 The data provider currently only supports reading fixed-size batches from input files, so data batches with different sizes will require some additional work.
 

psyche-book/src/development/python.md

@@ -1,13 +1,13 @@
 # Python Integration
 
 > [!WARNING]
-> Python support is still under development and not production-ready.
+> Python support is still under development.
 > The APIs used to write it are not documented
 > because they are still subject to large amounts of change.
 
 ## Overview
 
-Psyche provides a Python integration that allows you to write modeling code in Python using libraries like [Hugging Face Transformers](https://github.com/huggingface/transformers) while leveraging Psyche's Rust core for training orchestration. This integration is designed for research where you want the flexibility of Python modeling with Psyche's training infrastructure, and production-scale training where you want to take advantage of highly optimized training frameworks already built in Python.
+Psyche provides a Python integration that allows you to write modeling code in Python using libraries like [Hugging Face Transformers](https://github.com/huggingface/transformers) or [TorchTitan](https://github.com/pytorch/torchtitan) while leveraging Psyche's Rust core for training orchestration. This integration is designed for research where you want the flexibility of Python modeling with Psyche's training infrastructure, and production-scale training where you want to take advantage of highly optimized training frameworks already built in Python.
 
 The Python integration works through a "sidecar" process that Psyche spawns and communicates with during training.
 
@@ -20,6 +20,7 @@ This shell provides:
 - The `psyche` Python module (built from Rust using PyO3)
 - PyTorch
 - Transformers library
+- TorchTitan library
 - Other required Python dependencies via `pyproject.toml` / `uv.lock`
 
 ### Development Workflow
@@ -33,12 +34,11 @@ We recommend running commands directly through the dev shell without entering it
 For example, to run the `train` program using python:
 
 ```bash
-nix develop .#python --command just train-model-python \
+nix develop .#python --command cargo run --bin train \
   --model emozilla/llama2-20m-init \
   --data-path ./data/fineweb-10bt/ \
   --total-batch 2 \
-  --micro-batch 1 \
-  --python
+  --micro-batch 1
 ```
 
 Alternatively, you _could_ enter the shell and run the commands with:

psyche-book/src/enduser/run-config.md

@@ -64,8 +64,7 @@ total_steps = 25000
 ```toml
 # so far only LLMs are supported.
 [model.LLM]
-# Architecture of the model to train on can be HfLlama or HfDeepseek for now.
-# If running with Python sidecars this must be set to HfAuto.
+# Architecture of the model to train on can be HfLlama, HfDeepseek, or, when using the Python integration, HfAuto or TorchTitan.
 architecture = "HfLlama"
 data_type = "Pretraining"
 max_seq_len = 2048