feat(jepa): extended GPU training, data pipeline, integration, and cargo config

- jepa_gpu: remove the bring-up 2-block cap; full-depth GPU-resident
  ViT verified against a full-depth CPU reference on the RTX 5060 Ti
  (depth-12 ViT-Tiny max_rel_err <= 6.6e-5). CpuViTEncoder's own hidden
  min(depth,2) cap removed too — CPU-path callers now get the model
  they configured.
- jepa_distributed: real TCP parameter-server AllReduce backend
  (rendezvous handshake with world-size/rank validation, length-
  prefixed f32 payloads, connect/read/accept timeouts, connect retry
  until deadline so early peers survive rank 0 still computing);
  jepa_runner wires it for world_size > 1 and fails hard on collective
  errors. Two-rank loopback training run covered by test.
- rtx-jepa-cli (new crate): rtx-jepa binary with train/bench/plan/
  validate subcommands driving JepaRunConfig, run_jepa_training,
  run_jepa_benchmark, and ClusterTrainingPlan (plan --emit-config
  round-trips through a config serializer). GPU bench on this node:
  86k patches/sec vs 1.3k CPU (~64x).
- ViTSizeStr::Micro (d=32, depth=2) added as an explicit test/smoke
  size now that no hidden caps keep full-size configs cheap; heavy
  tests moved onto it (rtx-transformers suite: 367s -> 5s, and the
  runner subset had ballooned to 35min at full depth before this).

966 lib tests pass; 35/35 jepa_gpu with cuda; 8/8 CLI tests.

Co-Authored-By: Claude Fable 5 <[email protected]>

Cargo.toml

@@ -142,6 +142,7 @@ members = [
     # Development & tooling (13 crates)
     "crates/tooling/rtx-eval",
     "crates/tooling/rtx-bench",
+    "crates/tooling/rtx-jepa-cli",
     "crates/tooling/rtx-kernel-bench",
     
     # Data management (3 crates)

crates/tooling/rtx-jepa-cli/Cargo.toml

@@ -0,0 +1,29 @@
+[package]
+name = "rtx-jepa-cli"
+version = "1.0.0"
+edition.workspace = true
+rust-version = "1.92"
+authors = ["RustyTorch Team"]
+license = "MIT OR Apache-2.0"
+repository = "https://github.com/rustytorch/rustytorch"
+description = "CLI driver for JEPA (I-JEPA/V-JEPA) training, benchmarking, and cluster planning"
+keywords = ["jepa", "training", "cli", "ssl", "gpu"]
+categories = ["command-line-utilities", "science"]
+
+[dependencies]
+rtx-transformers = { workspace = true }
+clap = { workspace = true }
+anyhow = { workspace = true }
+
+[dev-dependencies]
+tempfile = { workspace = true }
+
+[features]
+cuda = ["rtx-transformers/cuda"]
+
+[[bin]]
+name = "rtx-jepa"
+path = "src/main.rs"
+
+[lints]
+workspace = true

crates/tooling/rtx-jepa-cli/src/config_fmt.rs

@@ -0,0 +1,134 @@
+//! Serializes a [`JepaRunConfig`] back to the `key = value` text format that
+//! `rtx_transformers::ssl::parse_config_from_str` understands.
+//!
+//! `rtx-transformers` only ships a parser (`parse_config_from_str`), not a
+//! serializer, so this lives here in the CLI crate rather than upstream.
+
+use rtx_transformers::ssl::{JepaRunConfig, ViTSizeStr};
+
+fn vit_size_str(v: &ViTSizeStr) -> &'static str {
+    match v {
+        ViTSizeStr::Micro => "micro",
+        ViTSizeStr::Tiny => "tiny",
+        ViTSizeStr::Small => "small",
+        ViTSizeStr::Base => "base",
+        ViTSizeStr::Large => "large",
+        ViTSizeStr::Huge => "huge",
+    }
+}
+
+fn quote(s: &str) -> String {
+    format!("\"{s}\"")
+}
+
+fn string_array(items: &[String]) -> String {
+    let inner = items
+        .iter()
+        .map(|s| quote(s))
+        .collect::<Vec<_>>()
+        .join(", ");
+    format!("[{inner}]")
+}
+
+/// Render a [`JepaRunConfig`] as a `key = value` config file, in the same
+/// key set that [`rtx_transformers::ssl::parse_config_from_str`] accepts.
+pub fn to_config_string(cfg: &JepaRunConfig) -> String {
+    let mut out = String::new();
+
+    out.push_str("# --- Model ---\n");
+    out.push_str(&format!("vit_size = {}\n", quote(vit_size_str(&cfg.vit_size))));
+    out.push_str(&format!("image_size = {}\n", cfg.image_size));
+    out.push_str(&format!("patch_size = {}\n", cfg.patch_size));
+
+    out.push_str("\n# --- Training schedule ---\n");
+    out.push_str(&format!("total_steps = {}\n", cfg.total_steps));
+    out.push_str(&format!("warmup_steps = {}\n", cfg.warmup_steps));
+    out.push_str(&format!("base_lr = {}\n", cfg.base_lr));
+    out.push_str(&format!("weight_decay = {}\n", cfg.weight_decay));
+    out.push_str(&format!("ema_tau_start = {}\n", cfg.ema_tau_start));
+    out.push_str(&format!("ema_tau_end = {}\n", cfg.ema_tau_end));
+
+    out.push_str("\n# --- Data ---\n");
+    out.push_str(&format!("batch_size = {}\n", cfg.batch_size));
+    out.push_str(&format!("num_workers = {}\n", cfg.num_workers));
+    out.push_str(&format!("data_shards = {}\n", string_array(&cfg.data_shards)));
+
+    out.push_str("\n# --- Checkpointing ---\n");
+    out.push_str(&format!("checkpoint_dir = {}\n", quote(&cfg.checkpoint_dir)));
+    out.push_str(&format!("checkpoint_every = {}\n", cfg.checkpoint_every));
+    out.push_str(&format!(
+        "resume_from = {}\n",
+        quote(cfg.resume_from.as_deref().unwrap_or(""))
+    ));
+
+    out.push_str("\n# --- Logging ---\n");
+    out.push_str(&format!("log_every = {}\n", cfg.log_every));
+    out.push_str(&format!("eval_every = {}\n", cfg.eval_every));
+
+    out.push_str("\n# --- Cluster ---\n");
+    out.push_str(&format!("num_gpus = {}\n", cfg.num_gpus));
+    out.push_str(&format!("tensor_parallel = {}\n", cfg.tensor_parallel));
+    out.push_str(&format!("data_parallel = {}\n", cfg.data_parallel));
+
+    out.push_str("\n# --- Distributed training ---\n");
+    out.push_str(&format!("world_size = {}\n", cfg.world_size));
+    out.push_str(&format!("rank = {}\n", cfg.rank));
+    out.push_str(&format!("master_addr = {}\n", quote(&cfg.master_addr)));
+    out.push_str(&format!("master_port = {}\n", cfg.master_port));
+
+    out.push_str("\n# --- Metrics export ---\n");
+    out.push_str(&format!(
+        "metrics_csv_path = {}\n",
+        quote(cfg.metrics_csv_path.as_deref().unwrap_or(""))
+    ));
+
+    out.push_str("\n# --- Benchmark mode ---\n");
+    out.push_str(&format!("benchmark_mode = {}\n", cfg.benchmark_mode));
+    out.push_str(&format!("benchmark_steps = {}\n", cfg.benchmark_steps));
+    out.push_str(&format!("benchmark_patches = {}\n", cfg.benchmark_patches));
+
+    out.push_str("\n# --- GPU encoder ---\n");
+    out.push_str(&format!("use_gpu = {}\n", cfg.use_gpu));
+    out.push_str(&format!("gpu_device_id = {}\n", cfg.gpu_device_id));
+
+    out
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use rtx_transformers::ssl::parse_config_from_str;
+
+    #[test]
+    fn default_config_round_trips() {
+        let cfg = JepaRunConfig::default();
+        let text = to_config_string(&cfg);
+        let parsed = parse_config_from_str(&text).unwrap();
+        assert!(matches!(parsed.vit_size, ViTSizeStr::Tiny));
+        assert_eq!(parsed.image_size, cfg.image_size);
+        assert_eq!(parsed.patch_size, cfg.patch_size);
+        assert_eq!(parsed.total_steps, cfg.total_steps);
+        assert_eq!(parsed.base_lr, cfg.base_lr);
+        assert_eq!(parsed.data_shards, cfg.data_shards);
+        assert_eq!(parsed.master_port, cfg.master_port);
+        assert_eq!(parsed.use_gpu, cfg.use_gpu);
+    }
+
+    #[test]
+    fn non_default_fields_round_trip() {
+        let mut cfg = JepaRunConfig::default();
+        cfg.vit_size = ViTSizeStr::Large;
+        cfg.resume_from = Some("ckpt.jepa".to_string());
+        cfg.metrics_csv_path = Some("metrics.csv".to_string());
+        cfg.data_shards = vec!["shard0.tar".to_string(), "shard1.tar".to_string()];
+        cfg.use_gpu = true;
+
+        let text = to_config_string(&cfg);
+        let parsed = parse_config_from_str(&text).unwrap();
+        assert!(matches!(parsed.vit_size, ViTSizeStr::Large));
+        assert_eq!(parsed.resume_from.as_deref(), Some("ckpt.jepa"));
+        assert_eq!(parsed.metrics_csv_path.as_deref(), Some("metrics.csv"));
+        assert_eq!(parsed.data_shards, vec!["shard0.tar", "shard1.tar"]);
+        assert!(parsed.use_gpu);
+    }
+}

crates/tooling/rtx-jepa-cli/src/main.rs

@@ -0,0 +1,399 @@
+//! `rtx-jepa` — CLI driver for the JEPA training runner in `rtx-transformers`.
+//!
+//! Subcommands:
+//! - `train`    — parse a config file and run a full training session
+//! - `bench`    — run the encoder throughput benchmark
+//! - `plan`     — build a cluster topology + parallelism plan and print it
+//! - `validate` — parse a config file and print the resolved config (or error)
+
+use std::path::PathBuf;
+
+use anyhow::{Context, Result};
+use clap::{Parser, Subcommand};
+
+use rtx_transformers::ssl::{
+    parse_config_from_str, run_jepa_training, ClusterTopology, JepaClusterConfig,
+    JepaParallelConfig, JepaRunConfig, ClusterTrainingPlan,
+};
+use rtx_transformers::ssl::jepa_runner::run_jepa_benchmark;
+
+mod config_fmt;
+use config_fmt::to_config_string;
+
+#[derive(Parser)]
+#[command(name = "rtx-jepa", version, about = "JEPA training CLI for RustyTorch")]
+struct Cli {
+    #[command(subcommand)]
+    command: Command,
+}
+
+#[derive(Subcommand)]
+enum Command {
+    /// Run a full JEPA training session from a config file.
+    Train {
+        /// Path to a `key = value` JEPA config file.
+        #[arg(long)]
+        config: PathBuf,
+        /// Force GPU encoder (overrides `use_gpu` in the config file).
+        #[arg(long)]
+        gpu: bool,
+        /// Resume from a checkpoint path (overrides `resume_from` in the config file).
+        #[arg(long)]
+        resume: Option<String>,
+        /// Write per-run metrics to this CSV path (overrides `metrics_csv_path`).
+        #[arg(long, value_name = "PATH")]
+        metrics_csv: Option<String>,
+    },
+    /// Run the encoder throughput benchmark (no real training).
+    Bench {
+        /// ViT backbone size.
+        #[arg(long, value_enum, default_value = "tiny")]
+        vit: VitSizeArg,
+        /// Number of benchmark steps to run.
+        #[arg(long, default_value_t = 50)]
+        steps: usize,
+        /// Use the GPU encoder (requires the `cuda` feature).
+        #[arg(long)]
+        gpu: bool,
+        /// Optional path to write benchmark results as CSV.
+        #[arg(long, value_name = "PATH")]
+        metrics_csv: Option<String>,
+    },
+    /// Build a cluster topology + parallelism plan and print it.
+    Plan {
+        /// ViT backbone size (drives parameter count for TP/PP selection).
+        #[arg(long, value_enum, default_value = "base")]
+        vit: VitSizeArg,
+        /// Number of nodes in the cluster.
+        #[arg(long, default_value_t = 1)]
+        nodes: usize,
+        /// GPUs per node.
+        #[arg(long = "gpus-per-node", default_value_t = 1)]
+        gpus_per_node: usize,
+        /// Dataset size in number of images/samples.
+        #[arg(long = "dataset-size", default_value_t = 1_281_167)]
+        dataset_size: usize,
+        /// Target number of training epochs.
+        #[arg(long, default_value_t = 1)]
+        epochs: usize,
+        /// Also print the resulting JepaRunConfig (after apply_cluster_plan)
+        /// as a config file the `train` command can consume.
+        #[arg(long = "emit-config")]
+        emit_config: bool,
+    },
+    /// Parse a config file and print the resolved config (or the parse error).
+    Validate {
+        /// Path to a `key = value` JEPA config file.
+        #[arg(long)]
+        config: PathBuf,
+    },
+}
+
+#[derive(Copy, Clone, Debug, clap::ValueEnum)]
+enum VitSizeArg {
+    Tiny,
+    Small,
+    Base,
+    Large,
+    Huge,
+}
+
+impl VitSizeArg {
+    fn as_str(self) -> &'static str {
+        match self {
+            VitSizeArg::Tiny => "tiny",
+            VitSizeArg::Small => "small",
+            VitSizeArg::Base => "base",
+            VitSizeArg::Large => "large",
+            VitSizeArg::Huge => "huge",
+        }
+    }
+
+    fn approx_params_m(self) -> usize {
+        match self {
+            VitSizeArg::Tiny => 6,
+            VitSizeArg::Small => 22,
+            VitSizeArg::Base => 86,
+            VitSizeArg::Large => 307,
+            VitSizeArg::Huge => 632,
+        }
+    }
+}
+
+fn main() -> Result<()> {
+    let cli = Cli::parse();
+
+    match cli.command {
+        Command::Train {
+            config,
+            gpu,
+            resume,
+            metrics_csv,
+        } => cmd_train(&config, gpu, resume, metrics_csv),
+        Command::Bench {
+            vit,
+            steps,
+            gpu,
+            metrics_csv,
+        } => cmd_bench(vit, steps, gpu, metrics_csv),
+        Command::Plan {
+            vit,
+            nodes,
+            gpus_per_node,
+            dataset_size,
+            epochs,
+            emit_config,
+        } => cmd_plan(vit, nodes, gpus_per_node, dataset_size, epochs, emit_config),
+        Command::Validate { config } => cmd_validate(&config),
+    }
+}
+
+fn cmd_train(
+    config_path: &PathBuf,
+    gpu: bool,
+    resume: Option<String>,
+    metrics_csv: Option<String>,
+) -> Result<()> {
+    let text = std::fs::read_to_string(config_path)
+        .with_context(|| format!("reading config file {}", config_path.display()))?;
+    let mut cfg = parse_config_from_str(&text)
+        .map_err(|e| anyhow::anyhow!("failed to parse config {}: {e}", config_path.display()))?;
+
+    if gpu {
+        cfg.use_gpu = true;
+    }
+    if let Some(resume) = resume {
+        cfg.resume_from = Some(resume);
+    }
+    if let Some(csv) = metrics_csv {
+        cfg.metrics_csv_path = Some(csv);
+    }
+
+    println!(
+        "Starting JEPA training: ViT-{:?} steps={} batch={} gpu={}",
+        cfg.vit_size, cfg.total_steps, cfg.batch_size, cfg.use_gpu
+    );
+
+    let summary = run_jepa_training(cfg);
+    print_training_summary(&summary);
+
+    Ok(())
+}
+
+fn print_training_summary(summary: &rtx_transformers::ssl::JepaTrainingSummary) {
+    println!("=== JEPA Training Summary ===");
+    println!("Total steps:          {}", summary.total_steps);
+    println!("Final loss:           {:.6}", summary.final_loss);
+    println!("Mean loss (last 100): {:.6}", summary.mean_loss);
+    println!("Steps/sec:            {:.2}", summary.steps_per_second);
+    println!("Tokens/sec:           {:.1}", summary.tokens_per_sec);
+    println!("Checkpoints saved:    {}", summary.checkpoints_saved);
+    println!("Wall time (s):        {:.2}", summary.wall_time_seconds);
+    println!("World size:           {}", summary.world_size);
+    println!("Effective batch:      {}", summary.effective_batch_size);
+    println!("Real-data steps:      {}", summary.real_data_steps);
+}
+
+fn cmd_bench(vit: VitSizeArg, steps: usize, gpu: bool, metrics_csv: Option<String>) -> Result<()> {
+    let mut cfg = JepaRunConfig::default();
+    cfg.vit_size = parse_vit_size(vit.as_str())?;
+    cfg.benchmark_steps = steps;
+    cfg.use_gpu = gpu;
+
+    let result = run_jepa_benchmark(&cfg);
+    result.print_summary();
+
+    if let Some(path) = metrics_csv {
+        std::fs::write(&path, result.to_csv())
+            .with_context(|| format!("writing benchmark CSV to {path}"))?;
+        println!("Wrote CSV to {path}");
+    }
+
+    Ok(())
+}
+
+fn cmd_plan(
+    vit: VitSizeArg,
+    nodes: usize,
+    gpus_per_node: usize,
+    dataset_size: usize,
+    epochs: usize,
+    emit_config: bool,
+) -> Result<()> {
+    let topology = ClusterTopology::rtx5060ti_cluster(nodes, gpus_per_node);
+    let global_batch_size = topology.total_gpus().max(1) * 64;
+    let parallel =
+        JepaParallelConfig::for_model_and_cluster(vit.approx_params_m(), &topology, global_batch_size);
+
+    let cluster_config = JepaClusterConfig {
+        topology,
+        parallel,
+        grad_compression: Default::default(),
+        checkpoint: Default::default(),
+        global_batch_size,
+        micro_batch_size: 64,
+        gradient_accumulation_steps: 1,
+        mixed_precision: rtx_transformers::ssl::MixedPrecision::Bf16,
+        compile_model: false,
+    };
+
+    let model_name = format!("JEPA-ViT-{}", vit.as_str());
+    let plan = ClusterTrainingPlan::new(cluster_config, &model_name, dataset_size, epochs);
+
+    println!("{}", plan.summary());
+
+    if emit_config {
+        let mut run_cfg = JepaRunConfig::default();
+        run_cfg.vit_size = parse_vit_size(vit.as_str())?;
+        let _ = run_cfg.apply_cluster_plan(&plan);
+
+        println!("\n=== Emitted config (train --config <this>) ===");
+        println!("{}", to_config_string(&run_cfg));
+    }
+
+    Ok(())
+}
+
+fn cmd_validate(config_path: &PathBuf) -> Result<()> {
+    let text = std::fs::read_to_string(config_path)
+        .with_context(|| format!("reading config file {}", config_path.display()))?;
+
+    match parse_config_from_str(&text) {
+        Ok(cfg) => {
+            println!("Config OK: {}", config_path.display());
+            println!("{}", to_config_string(&cfg));
+            Ok(())
+        }
+        Err(e) => {
+            eprintln!("Config error in {}: {e}", config_path.display());
+            std::process::exit(1);
+        }
+    }
+}
+
+fn parse_vit_size(s: &str) -> Result<rtx_transformers::ssl::ViTSizeStr> {
+    use rtx_transformers::ssl::ViTSizeStr;
+    match s {
+        "micro" => Ok(ViTSizeStr::Micro),
+        "tiny" => Ok(ViTSizeStr::Tiny),
+        "small" => Ok(ViTSizeStr::Small),
+        "base" => Ok(ViTSizeStr::Base),
+        "large" => Ok(ViTSizeStr::Large),
+        "huge" => Ok(ViTSizeStr::Huge),
+        other => anyhow::bail!("unknown vit size: {other}"),
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use clap::CommandFactory;
+
+    #[test]
+    fn cli_verifies() {
+        Cli::command().debug_assert();
+    }
+
+    #[test]
+    fn parses_train_args() {
+        let cli = Cli::try_parse_from([
+            "rtx-jepa",
+            "train",
+            "--config",
+            "run.toml",
+            "--gpu",
+            "--resume",
+            "ckpt.jepa",
+            "--metrics-csv",
+            "out.csv",
+        ])
+        .unwrap();
+        match cli.command {
+            Command::Train {
+                config,
+                gpu,
+                resume,
+                metrics_csv,
+            } => {
+                assert_eq!(config, PathBuf::from("run.toml"));
+                assert!(gpu);
+                assert_eq!(resume.as_deref(), Some("ckpt.jepa"));
+                assert_eq!(metrics_csv.as_deref(), Some("out.csv"));
+            }
+            _ => panic!("expected Train"),
+        }
+    }
+
+    #[test]
+    fn parses_bench_args() {
+        let cli = Cli::try_parse_from([
+            "rtx-jepa", "bench", "--vit", "small", "--steps", "10", "--gpu",
+        ])
+        .unwrap();
+        match cli.command {
+            Command::Bench { vit, steps, gpu, .. } => {
+                assert!(matches!(vit, VitSizeArg::Small));
+                assert_eq!(steps, 10);
+                assert!(gpu);
+            }
+            _ => panic!("expected Bench"),
+        }
+    }
+
+    #[test]
+    fn parses_plan_args() {
+        let cli = Cli::try_parse_from([
+            "rtx-jepa",
+            "plan",
+            "--vit",
+            "large",
+            "--nodes",
+            "2",
+            "--gpus-per-node",
+            "4",
+            "--dataset-size",
+            "1000000",
+            "--epochs",
+            "5",
+            "--emit-config",
+        ])
+        .unwrap();
+        match cli.command {
+            Command::Plan {
+                vit,
+                nodes,
+                gpus_per_node,
+                dataset_size,
+                epochs,
+                emit_config,
+            } => {
+                assert!(matches!(vit, VitSizeArg::Large));
+                assert_eq!(nodes, 2);
+                assert_eq!(gpus_per_node, 4);
+                assert_eq!(dataset_size, 1_000_000);
+                assert_eq!(epochs, 5);
+                assert!(emit_config);
+            }
+            _ => panic!("expected Plan"),
+        }
+    }
+
+    #[test]
+    fn parses_validate_args() {
+        let cli = Cli::try_parse_from(["rtx-jepa", "validate", "--config", "run.toml"]).unwrap();
+        match cli.command {
+            Command::Validate { config } => assert_eq!(config, PathBuf::from("run.toml")),
+            _ => panic!("expected Validate"),
+        }
+    }
+
+    #[test]
+    fn config_round_trip() {
+        let cfg = JepaRunConfig::default();
+        let text = to_config_string(&cfg);
+        let parsed = parse_config_from_str(&text).expect("round-trip parse should succeed");
+        assert_eq!(parsed.total_steps, cfg.total_steps);
+        assert_eq!(parsed.batch_size, cfg.batch_size);
+        assert_eq!(parsed.checkpoint_dir, cfg.checkpoint_dir);
+    }
+}

crates/training/rtx-transformers/src/ssl/jepa_distributed.rs

@@ -28,6 +28,308 @@ const ALLREDUCE_TIMEOUT_SECS: u64 = 1;
 #[cfg(not(test))]
 const ALLREDUCE_TIMEOUT_SECS: u64 = 30;
 
+/// Timeout for the unconditionally-compiled `Tcp` backend (connect, read,
+/// and accept-poll deadlines). Short in tests so a missing peer errors
+/// quickly instead of hanging the suite; longer in production so slow
+/// cluster rendezvous doesn't spuriously fail.
+#[cfg(test)]
+const TCP_TIMEOUT_SECS: u64 = 2;
+
+#[cfg(not(test))]
+const TCP_TIMEOUT_SECS: u64 = 30;
+
+/// Poll interval used while waiting on a non-blocking `accept()` for the
+/// timeout-capable listener helper below.
+const TCP_ACCEPT_POLL_MS: u64 = 10;
+
+// ============================================================================
+// TcpAllReduce — real, unconditionally-compiled parameter-server AllReduce
+// ============================================================================
+//
+// Protocol (single connection per non-primary rank, reused for both
+// directions — no second port needed):
+//
+// 1. Handshake: the connecting (non-zero) rank sends an 8-byte header
+//    `[rank: u32 LE][world_size: u32 LE]`. Rank 0 validates that the
+//    peer's `world_size` matches its own and that `rank` is a unique
+//    value in `1..world_size`. On success rank 0 replies with a single
+//    `0x01` ack byte; on mismatch it replies `0x00` followed by a
+//    length-prefixed UTF-8 error string, then drops the connection.
+// 2. Gradient send: the peer writes `[len: u64 LE][f32 LE payload]`.
+// 3. Averaged broadcast: rank 0 writes `[len: u64 LE][f32 LE payload]`
+//    back down the *same* socket once all peers have reported in.
+//
+// All reads/writes use `read_exact`/`write_all` (which internally loop
+// over partial reads/writes) and every socket has a read/connect timeout
+// so a missing or hung peer produces an `Err` instead of blocking
+// forever.
+mod tcp_allreduce {
+    use std::io::{Read, Write};
+    use std::net::{SocketAddr, TcpListener, TcpStream};
+    use std::time::{Duration, Instant};
+
+    use super::TCP_ACCEPT_POLL_MS;
+
+    fn read_u32(stream: &mut TcpStream) -> Result<u32, String> {
+        let mut buf = [0u8; 4];
+        stream
+            .read_exact(&mut buf)
+            .map_err(|e| format!("read u32 failed: {e}"))?;
+        Ok(u32::from_le_bytes(buf))
+    }
+
+    fn write_u32(stream: &mut TcpStream, v: u32) -> Result<(), String> {
+        stream
+            .write_all(&v.to_le_bytes())
+            .map_err(|e| format!("write u32 failed: {e}"))
+    }
+
+    fn read_len_prefixed(stream: &mut TcpStream) -> Result<Vec<u8>, String> {
+        let mut len_buf = [0u8; 8];
+        stream
+            .read_exact(&mut len_buf)
+            .map_err(|e| format!("read length prefix failed: {e}"))?;
+        let len = u64::from_le_bytes(len_buf) as usize;
+        let mut payload = vec![0u8; len];
+        stream
+            .read_exact(&mut payload)
+            .map_err(|e| format!("read payload ({len} bytes) failed: {e}"))?;
+        Ok(payload)
+    }
+
+    fn write_len_prefixed(stream: &mut TcpStream, payload: &[u8]) -> Result<(), String> {
+        stream
+            .write_all(&(payload.len() as u64).to_le_bytes())
+            .map_err(|e| format!("write length prefix failed: {e}"))?;
+        stream
+            .write_all(payload)
+            .map_err(|e| format!("write payload failed: {e}"))
+    }
+
+    fn f32_to_bytes(grads: &[f32]) -> Vec<u8> {
+        grads.iter().flat_map(|v| v.to_le_bytes()).collect()
+    }
+
+    fn bytes_to_f32(bytes: &[u8]) -> Result<Vec<f32>, String> {
+        if bytes.len() % 4 != 0 {
+            return Err(format!(
+                "gradient payload length {} is not a multiple of 4",
+                bytes.len()
+            ));
+        }
+        Ok(bytes
+            .chunks_exact(4)
+            .map(|c| f32::from_le_bytes(c.try_into().unwrap()))
+            .collect())
+    }
+
+    /// Accept a connection with an overall deadline by polling a
+    /// non-blocking listener. Returns `Err` instead of hanging forever
+    /// when no peer ever connects.
+    fn accept_with_timeout(
+        listener: &TcpListener,
+        timeout: Duration,
+    ) -> Result<(TcpStream, SocketAddr), String> {
+        listener
+            .set_nonblocking(true)
+            .map_err(|e| format!("set_nonblocking failed: {e}"))?;
+        let deadline = Instant::now() + timeout;
+        loop {
+            match listener.accept() {
+                Ok((stream, addr)) => {
+                    stream.set_nonblocking(false).ok();
+                    return Ok((stream, addr));
+                }
+                Err(e) if e.kind() == std::io::ErrorKind::WouldBlock => {
+                    if Instant::now() >= deadline {
+                        return Err(format!(
+                            "accept timed out after {:?} waiting for a peer",
+                            timeout
+                        ));
+                    }
+                    std::thread::sleep(Duration::from_millis(TCP_ACCEPT_POLL_MS));
+                }
+                Err(e) => return Err(format!("accept failed: {e}")),
+            }
+        }
+    }
+
+    /// Rank 0: gather grads from every other rank, sum + average, then
+    /// broadcast the result back down each connection.
+    fn run_primary(
+        world_size: usize,
+        master_addr: &str,
+        master_port: u16,
+        grads: &mut Vec<f32>,
+        timeout: Duration,
+    ) -> Result<(), String> {
+        let bind_addr = format!("{master_addr}:{master_port}");
+        let listener = TcpListener::bind(&bind_addr)
+            .map_err(|e| format!("rank 0 bind failed on {bind_addr}: {e}"))?;
+
+        let mut sum = grads.clone();
+        let mut seen_ranks: std::collections::HashSet<u32> = std::collections::HashSet::new();
+        let mut peers: Vec<TcpStream> = Vec::with_capacity(world_size - 1);
+
+        while peers.len() < world_size - 1 {
+            let (mut stream, peer_addr) = accept_with_timeout(&listener, timeout)?;
+            stream.set_read_timeout(Some(timeout)).ok();
+            stream.set_write_timeout(Some(timeout)).ok();
+
+            // --- handshake ---
+            let peer_rank = read_u32(&mut stream)
+                .map_err(|e| format!("handshake read rank from {peer_addr} failed: {e}"))?;
+            let peer_world_size = read_u32(&mut stream)
+                .map_err(|e| format!("handshake read world_size from {peer_addr} failed: {e}"))?;
+
+            let mismatch = if peer_world_size as usize != world_size {
+                Some(format!(
+                    "world_size mismatch: rank 0 expects {world_size}, peer {peer_addr} (rank {peer_rank}) sent {peer_world_size}"
+                ))
+            } else if peer_rank == 0 || peer_rank as usize >= world_size {
+                Some(format!(
+                    "invalid rank {peer_rank} from {peer_addr} for world_size {world_size}"
+                ))
+            } else if !seen_ranks.insert(peer_rank) {
+                Some(format!("duplicate rank {peer_rank} from {peer_addr}"))
+            } else {
+                None
+            };
+
+            if let Some(reason) = mismatch {
+                // Reject: ack=0 then length-prefixed reason, then drop.
+                stream.write_all(&[0u8]).ok();
+                write_len_prefixed(&mut stream, reason.as_bytes()).ok();
+                return Err(reason);
+            }
+
+            // Accept: ack=1
+            stream
+                .write_all(&[1u8])
+                .map_err(|e| format!("handshake ack write to {peer_addr} failed: {e}"))?;
+
+            // --- receive this peer's gradients, accumulate ---
+            let payload = read_len_prefixed(&mut stream)
+                .map_err(|e| format!("grad recv from {peer_addr} failed: {e}"))?;
+            let peer_grads = bytes_to_f32(&payload)?;
+            if peer_grads.len() != sum.len() {
+                return Err(format!(
+                    "grad length mismatch from {peer_addr}: expected {}, got {}",
+                    sum.len(),
+                    peer_grads.len()
+                ));
+            }
+            for (s, v) in sum.iter_mut().zip(peer_grads.iter()) {
+                *s += v;
+            }
+
+            peers.push(stream);
+        }
+
+        for v in &mut sum {
+            *v /= world_size as f32;
+        }
+        let avg_bytes = f32_to_bytes(&sum);
+
+        for mut stream in peers {
+            write_len_prefixed(&mut stream, &avg_bytes)
+                .map_err(|e| format!("broadcast send failed: {e}"))?;
+        }
+
+        *grads = sum;
+        Ok(())
+    }
+
+    /// Non-zero rank: connect, handshake, send local grads, receive the
+    /// averaged result.
+    fn run_peer(
+        world_size: usize,
+        rank: usize,
+        master_addr: &str,
+        master_port: u16,
+        grads: &mut Vec<f32>,
+        timeout: Duration,
+    ) -> Result<(), String> {
+        let addr = format!("{master_addr}:{master_port}");
+        let sock_addr: SocketAddr = addr
+            .parse()
+            .map_err(|e| format!("parse master addr '{addr}': {e}"))?;
+        // Rank 0 only opens its listener once it reaches its own all_reduce
+        // call, so peers arriving early see "connection refused". Retry with
+        // a short backoff until the shared deadline instead of failing on
+        // the first attempt.
+        let deadline = std::time::Instant::now() + timeout;
+        let mut stream = loop {
+            match TcpStream::connect_timeout(&sock_addr, timeout) {
+                Ok(s) => break s,
+                Err(e) => {
+                    if std::time::Instant::now() >= deadline {
+                        return Err(format!(
+                            "rank {rank} connect to {addr} failed after retrying for {timeout:?}: {e}"
+                        ));
+                    }
+                    std::thread::sleep(Duration::from_millis(50));
+                }
+            }
+        };
+        stream.set_read_timeout(Some(timeout)).ok();
+        stream.set_write_timeout(Some(timeout)).ok();
+
+        // --- handshake ---
+        write_u32(&mut stream, rank as u32)
+            .map_err(|e| format!("handshake send rank failed: {e}"))?;
+        write_u32(&mut stream, world_size as u32)
+            .map_err(|e| format!("handshake send world_size failed: {e}"))?;
+
+        let mut ack = [0u8; 1];
+        stream
+            .read_exact(&mut ack)
+            .map_err(|e| format!("handshake ack read failed: {e}"))?;
+        if ack[0] != 1 {
+            let reason_bytes = read_len_prefixed(&mut stream)
+                .unwrap_or_else(|_| b"rendezvous rejected".to_vec());
+            let reason = String::from_utf8_lossy(&reason_bytes).into_owned();
+            return Err(format!("rank {rank} rejected by rank 0: {reason}"));
+        }
+
+        // --- send local grads ---
+        let bytes = f32_to_bytes(grads);
+        write_len_prefixed(&mut stream, &bytes)
+            .map_err(|e| format!("grad send failed: {e}"))?;
+
+        // --- receive averaged grads ---
+        let payload =
+            read_len_prefixed(&mut stream).map_err(|e| format!("avg recv failed: {e}"))?;
+        let avg = bytes_to_f32(&payload)?;
+        if avg.len() != grads.len() {
+            return Err(format!(
+                "averaged grad length mismatch: expected {}, got {}",
+                grads.len(),
+                avg.len()
+            ));
+        }
+        *grads = avg;
+        Ok(())
+    }
+
+    /// Entry point: dispatches to `run_primary` or `run_peer` based on rank.
+    pub(super) fn all_reduce(
+        world_size: usize,
+        rank: usize,
+        master_addr: &str,
+        master_port: u16,
+        grads: &mut Vec<f32>,
+        timeout_secs: u64,
+    ) -> Result<(), String> {
+        let timeout = Duration::from_secs(timeout_secs);
+        if rank == 0 {
+            run_primary(world_size, master_addr, master_port, grads, timeout)
+        } else {
+            run_peer(world_size, rank, master_addr, master_port, grads, timeout)
+        }
+    }
+}
+
 // ============================================================================
 // TcpRingAllReduce  (only compiled when distributed-tcp feature is active)
 // ============================================================================
@@ -172,6 +474,11 @@ pub enum GradSyncBackend {
     /// `distributed-tcp` feature is enabled; falls back to simulated divide
     /// on any network error.
     Nccl { master_addr: String, master_port: u16 },
+    /// Real TCP parameter-server AllReduce (always compiled, no feature
+    /// gate, no external deps — plain `std::net`). Unlike `Nccl`, this
+    /// backend performs a rendezvous handshake that validates world_size
+    /// and rank before exchanging any gradient data.
+    Tcp { master_addr: String, master_port: u16 },
 }
 
 // ============================================================================
@@ -241,6 +548,75 @@ impl JepaGradSync {
         })
     }
 
+    /// Real TCP AllReduce config (validates world_size > 0, rank < world_size,
+    /// and port > 0). Unlike [`Self::nccl`], gradient synchronization over
+    /// this backend performs an actual rendezvous handshake + parameter-
+    /// server-style sum/average/broadcast using plain `std::net` sockets —
+    /// no external dependencies, no feature flag required.
+    pub fn tcp(
+        world_size: usize,
+        rank: usize,
+        master_addr: &str,
+        master_port: u16,
+    ) -> Result<Self, String> {
+        if world_size == 0 {
+            return Err("world_size must be > 0".to_string());
+        }
+        if rank >= world_size {
+            return Err(format!(
+                "rank {rank} must be < world_size {world_size}"
+            ));
+        }
+        if master_port == 0 {
+            return Err("master_port must be > 0".to_string());
+        }
+        Ok(Self {
+            world_size,
+            rank,
+            backend: GradSyncBackend::Tcp {
+                master_addr: master_addr.to_string(),
+                master_port,
+            },
+        })
+    }
+
+    /// Perform a real AllReduce over the `Tcp` backend, returning `Err` on
+    /// any network/protocol failure instead of silently falling back to
+    /// simulated averaging. No-op (`Ok`) for `world_size == 1`. For any
+    /// other backend this simply calls [`Self::sync_gradients`] and wraps
+    /// the result in `Ok`.
+    pub fn all_reduce(
+        &self,
+        grads: &mut Vec<f32>,
+        batch_local_loss: f32,
+    ) -> Result<GradSyncResult, String> {
+        if self.world_size <= 1 {
+            return Ok(self.sync_gradients(grads, batch_local_loss));
+        }
+
+        if let GradSyncBackend::Tcp { master_addr, master_port } = &self.backend {
+            tcp_allreduce::all_reduce(
+                self.world_size,
+                self.rank,
+                master_addr,
+                *master_port,
+                grads,
+                TCP_TIMEOUT_SECS,
+            )?;
+            let effective_batch_size = self.effective_batch_size(1);
+            let comm_bytes =
+                (grads.len() * std::mem::size_of::<f32>()) as u64 * self.world_size as u64;
+            return Ok(GradSyncResult {
+                world_size: self.world_size,
+                effective_batch_size,
+                loss_scale: 1.0f32 / self.world_size as f32,
+                comm_bytes,
+            });
+        }
+
+        Ok(self.sync_gradients(grads, batch_local_loss))
+    }
+
     /// Synchronize gradients across world_size processes.
     ///
     /// - `world_size == 1`: no-op, returns identity result
@@ -313,6 +689,34 @@ impl JepaGradSync {
                     comm_bytes,
                 }
             }
+            GradSyncBackend::Tcp { master_addr, master_port } => {
+                let mut comm_bytes = 0u64;
+                if let Err(e) = tcp_allreduce::all_reduce(
+                    self.world_size,
+                    self.rank,
+                    master_addr,
+                    *master_port,
+                    grads,
+                    TCP_TIMEOUT_SECS,
+                ) {
+                    // Network/protocol failure: log and fall back to
+                    // simulated averaging so the training loop is never
+                    // hard-blocked (same policy as the `Nccl` backend).
+                    eprintln!("TCP allreduce failed (falling back to simulated): {e}");
+                    for g in grads.iter_mut() {
+                        *g /= self.world_size as f32;
+                    }
+                } else {
+                    comm_bytes = (grads.len() * std::mem::size_of::<f32>()) as u64
+                        * self.world_size as u64;
+                }
+                GradSyncResult {
+                    world_size: self.world_size,
+                    effective_batch_size,
+                    loss_scale,
+                    comm_bytes,
+                }
+            }
         }
     }
 
@@ -335,6 +739,16 @@ impl JepaGradSync {
                 // Stub: real barrier would synchronize all ranks via NCCL
                 Ok(())
             }
+            GradSyncBackend::Tcp { .. } => {
+                // A zero-length allreduce round-trips through every rank,
+                // which is sufficient as a barrier.
+                if self.world_size <= 1 {
+                    Ok(())
+                } else {
+                    let mut empty: Vec<f32> = Vec::new();
+                    self.all_reduce(&mut empty, 0.0).map(|_| ())
+                }
+            }
         }
     }
 }
@@ -493,28 +907,40 @@ mod tests {
         assert!(sync.barrier().is_ok(), "barrier should return Ok");
     }
 
-    // 15. run_jepa_training with world_size=2, rank=0, total_steps=5
+    // 15. run_jepa_training with world_size=2: a real two-rank training run
+    // over the loopback TCP AllReduce backend (both ranks in-process).
     #[test]
     fn test_run_training_with_world_size_2() {
         use crate::ssl::jepa_runner::{JepaRunConfig, ViTSizeStr, run_jepa_training};
 
         let batch_size = 4usize;
-        let config = JepaRunConfig {
+        let make_config = |rank: usize| JepaRunConfig {
-            vit_size: ViTSizeStr::Tiny,
+            vit_size: ViTSizeStr::Micro,
             total_steps: 5,
             warmup_steps: 1,
             log_every: 100,           // suppress output
             checkpoint_every: 10_000, // no checkpoint
             batch_size,
             world_size: 2,
-            rank: 0,
+            rank,
             master_addr: "127.0.0.1".to_string(),
-            master_port: 29500,
+            // Dedicated port: avoid collisions with the tcp_allreduce unit
+            // tests (30001-30040) and any default-29500 usage.
+            master_port: 30050,
             ..JepaRunConfig::default()
         };
-        let summary = run_jepa_training(config);
+
+        let config0 = make_config(0);
+        let config1 = make_config(1);
+        let rank1 = std::thread::spawn(move || run_jepa_training(config1));
+        let summary0 = run_jepa_training(config0);
+        let summary1 = rank1.join().expect("rank 1 thread panicked");
+
+        for summary in [&summary0, &summary1] {
             assert_eq!(summary.world_size, 2);
             assert_eq!(summary.effective_batch_size, 2 * batch_size);
+            assert!(summary.final_loss.is_finite());
+        }
     }
 
     // -------------------------------------------------------------------------
@@ -581,4 +1007,151 @@ mod tests {
         let sync = JepaGradSync::nccl(2, 0, "127.0.0.1", 29503).unwrap();
         assert!(sync.barrier().is_ok());
     }
+
+    // -------------------------------------------------------------------------
+    // Real TCP AllReduce backend (`GradSyncBackend::Tcp`)
+    // -------------------------------------------------------------------------
+
+    // 21. tcp() constructor validation
+    #[test]
+    fn test_tcp_config_validation() {
+        assert!(JepaGradSync::tcp(2, 0, "127.0.0.1", 30001).is_ok());
+        assert!(
+            JepaGradSync::tcp(0, 0, "127.0.0.1", 30002).is_err(),
+            "world_size=0 should be rejected"
+        );
+        assert!(
+            JepaGradSync::tcp(2, 2, "127.0.0.1", 30003).is_err(),
+            "rank >= world_size should be rejected"
+        );
+        assert!(
+            JepaGradSync::tcp(2, 0, "127.0.0.1", 0).is_err(),
+            "port=0 should be rejected"
+        );
+    }
+
+    // 22. real 2-rank AllReduce over loopback TCP: exact expected mean
+    #[test]
+    fn test_tcp_allreduce_world_size_2() {
+        let addr = "127.0.0.1";
+        let port = 30010u16;
+
+        let sync1 = JepaGradSync::tcp(2, 1, addr, port).unwrap();
+        let peer = std::thread::spawn(move || {
+            let mut grads = vec![30.0f32, 40.0];
+            let result = sync1.all_reduce(&mut grads, 1.0);
+            (result, grads)
+        });
+
+        // Give rank 0 a head start binding the listener so rank 1's connect
+        // doesn't spuriously race an unbound socket (not required for
+        // correctness — the connect timeout would still recover — but it
+        // keeps the test fast and deterministic).
+        let sync0 = JepaGradSync::tcp(2, 0, addr, port).unwrap();
+        let mut grads0 = vec![10.0f32, 20.0];
+        let result0 = sync0.all_reduce(&mut grads0, 1.0);
+
+        let (result1, grads1) = peer.join().expect("peer thread panicked");
+
+        assert!(result0.is_ok(), "rank 0 all_reduce failed: {:?}", result0);
+        assert!(result1.is_ok(), "rank 1 all_reduce failed: {:?}", result1);
+        assert_eq!(grads0, vec![20.0f32, 30.0], "rank 0 mean mismatch");
+        assert_eq!(grads1, vec![20.0f32, 30.0], "rank 1 mean mismatch");
+
+        let r0 = result0.unwrap();
+        assert_eq!(r0.world_size, 2);
+        assert_eq!(r0.comm_bytes, 2 * 4 * 2); // 2 f32s * 4 bytes * world_size
+    }
+
+    // 23. real 3-rank AllReduce over loopback TCP: exact expected mean
+    #[test]
+    fn test_tcp_allreduce_world_size_3() {
+        let addr = "127.0.0.1";
+        let port = 30020u16;
+
+        let sync1 = JepaGradSync::tcp(3, 1, addr, port).unwrap();
+        let peer1 = std::thread::spawn(move || {
+            let mut grads = vec![2.0f32, 4.0];
+            let result = sync1.all_reduce(&mut grads, 1.0);
+            (result, grads)
+        });
+
+        let sync2 = JepaGradSync::tcp(3, 2, addr, port).unwrap();
+        let peer2 = std::thread::spawn(move || {
+            let mut grads = vec![3.0f32, 6.0];
+            let result = sync2.all_reduce(&mut grads, 1.0);
+            (result, grads)
+        });
+
+        let sync0 = JepaGradSync::tcp(3, 0, addr, port).unwrap();
+        let mut grads0 = vec![1.0f32, 2.0];
+        let result0 = sync0.all_reduce(&mut grads0, 1.0);
+
+        let (result1, grads1) = peer1.join().expect("peer1 thread panicked");
+        let (result2, grads2) = peer2.join().expect("peer2 thread panicked");
+
+        assert!(result0.is_ok(), "rank 0 all_reduce failed: {:?}", result0);
+        assert!(result1.is_ok(), "rank 1 all_reduce failed: {:?}", result1);
+        assert!(result2.is_ok(), "rank 2 all_reduce failed: {:?}", result2);
+
+        // sum = [6, 12], mean = [2, 4]
+        assert_eq!(grads0, vec![2.0f32, 4.0]);
+        assert_eq!(grads1, vec![2.0f32, 4.0]);
+        assert_eq!(grads2, vec![2.0f32, 4.0]);
+    }
+
+    // 24. mismatched world_size during handshake is rejected on both sides,
+    // not silently accepted.
+    #[test]
+    fn test_tcp_handshake_world_size_mismatch_rejected() {
+        let addr = "127.0.0.1";
+        let port = 30030u16;
+
+        // Rank 0 expects world_size=2, but the connecting peer claims 3.
+        let bad_peer = JepaGradSync::tcp(3, 1, addr, port).unwrap();
+        let peer = std::thread::spawn(move || {
+            let mut grads = vec![1.0f32];
+            bad_peer.all_reduce(&mut grads, 1.0)
+        });
+
+        let sync0 = JepaGradSync::tcp(2, 0, addr, port).unwrap();
+        let mut grads0 = vec![1.0f32];
+        let result0 = sync0.all_reduce(&mut grads0, 1.0);
+
+        let peer_result = peer.join().expect("peer thread panicked");
+
+        assert!(
+            result0.is_err(),
+            "rank 0 should reject a world_size-mismatched peer"
+        );
+        assert!(
+            peer_result.is_err(),
+            "mismatched peer should receive a rejection, not succeed"
+        );
+    }
+
+    // 25. a rank waiting for a peer that never connects times out with an
+    // Err instead of hanging forever.
+    #[test]
+    fn test_tcp_allreduce_timeout_no_peer() {
+        let addr = "127.0.0.1";
+        let port = 30040u16;
+
+        let sync0 = JepaGradSync::tcp(2, 0, addr, port).unwrap();
+        let mut grads0 = vec![1.0f32];
+
+        let start = std::time::Instant::now();
+        let result0 = sync0.all_reduce(&mut grads0, 1.0);
+        let elapsed = start.elapsed();
+
+        assert!(
+            result0.is_err(),
+            "all_reduce with no peer should error, not hang"
+        );
+        assert!(
+            elapsed < std::time::Duration::from_secs(10),
+            "all_reduce should time out promptly, took {:?}",
+            elapsed
+        );
+    }
 }

crates/training/rtx-transformers/src/ssl/jepa_gpu.rs

@@ -678,7 +678,7 @@ impl GpuViTEncoder {
         let total_us = total_start.elapsed().as_micros() as u64;
         let _patch_proj_elapsed = patch_proj_start.elapsed().as_micros() as u64;
 
-        let n_blocks = self.cpu_encoder.blocks.len().min(2);
+        let n_blocks = self.cpu_encoder.blocks.len();
 
         // Approximate split: proj ~25% of total, each block splits the rest evenly.
         let patch_proj_us = total_us / 4;
@@ -738,7 +738,10 @@ impl GpuViTEncoder {
         let d = self.cpu_encoder.config.embed_dim;
         let np = self.cpu_encoder.config.num_patches();
         let ffn_dim = (d as f32 * self.cpu_encoder.config.mlp_ratio) as usize;
-        let active_blocks = self.cpu_encoder.blocks.len().min(2);
+        // Run every configured block GPU-resident (no artificial depth cap —
+        // Batch 33: removed the bring-up-era `min(2)` cap now that per-block
+        // buffer lifetimes and VRAM usage have been verified at full depth).
+        let active_blocks = self.cpu_encoder.blocks.len();
 
         let kernels = self
             .kernels
@@ -1490,7 +1493,8 @@ mod tests {
     fn test_encode_with_timing_nonzero() {
         let enc = GpuViTEncoder::cpu(tiny_cfg());
         let (_, timing) = enc.encode_with_timing(&[0, 1, 2, 3, 4]);
-        assert_eq!(timing.block_us.len(), enc.cpu_encoder.blocks.len().min(2));
+        // No depth cap: one timing entry per configured block (tiny() → depth 12).
+        assert_eq!(timing.block_us.len(), enc.cpu_encoder.blocks.len());
     }
 
     // ── 22. tokens_per_sec with zero timing returns 0.0 ──────────────────────
@@ -1629,14 +1633,107 @@ mod tests {
         assert!((g1 - 0.841).abs() < 0.01, "GELU(1)≈0.841, got {g1}");
     }
 
-    // ── 34. Numerical parity: CpuViTEncoder vs GpuViTEncoder (cuda feature) ──
+    // ── Full-depth CPU reference (bypasses CpuViTEncoder::encode's own
+    // independent `min(depth, 2)` cap, which lives in jepa_vit.rs and is out
+    // of scope for this file) ─────────────────────────────────────────────
+    //
+    // `CpuViTEncoder::encode` (in jepa_vit.rs) still only runs the first two
+    // blocks — that cap is bring-up scaffolding in a sibling module this
+    // change does not touch. To get a genuine full-depth *reference* value
+    // to compare the now-uncapped GPU path against, this re-derives the same
+    // per-block math (layer norm → QKV → per-head SDPA → out-proj + residual
+    // → layer norm → FFN + residual) directly from `ViTBlock`'s public(crate)
+    // weight fields, using the exact same CPU helper functions
+    // (`layer_norm_rows_cpu`, `linear_cpu`, `sdp_attention_cpu`, `gelu_cpu`)
+    // already used elsewhere in this file, run over *all* configured blocks.
+    #[cfg(feature = "cuda")]
+    fn cpu_reference_full_depth(enc: &CpuViTEncoder, patch_indices: &[usize]) -> Vec<f32> {
+        let n = patch_indices.len();
+        if n == 0 {
+            return Vec::new();
+        }
+        let d = enc.config.embed_dim;
+        let np = enc.config.num_patches();
+
+        let mut tokens = vec![0.0f32; n * d];
+        for (i, &pi) in patch_indices.iter().enumerate() {
+            let pi = pi.min(np - 1);
+            tokens[i * d..(i + 1) * d].copy_from_slice(&enc.patch_embed[pi * d..(pi + 1) * d]);
+        }
+
+        let mut hidden = linear_cpu(&tokens, &enc.proj_w, &enc.proj_b, n, d, d);
+        for (i, &pi) in patch_indices.iter().enumerate() {
+            let pi = pi.min(np - 1);
+            for dd in 0..d {
+                hidden[i * d + dd] += enc.patch_embed[pi * d + dd] * 0.1;
+            }
+        }
+
+        for blk in &enc.blocks {
+            let h = blk.h;
+            let head_dim = d / h;
+            let ffn = blk.ffn1_b.len();
+
+            let mut normed = hidden.clone();
+            layer_norm_rows_cpu(&mut normed, d);
+            let qkv = linear_cpu(&normed, &blk.qkv_w, &blk.qkv_b, n, d, 3 * d);
+            let qkv_s: &[f32] = &qkv;
+
+            let mut attn_out = vec![0.0f32; n * d];
+            for hi in 0..h {
+                let q: Vec<f32> = (0..n)
+                    .flat_map(|t| (0..head_dim).map(move |dh| qkv_s[t * 3 * d + hi * head_dim + dh]))
+                    .collect();
+                let k: Vec<f32> = (0..n)
+                    .flat_map(|t| {
+                        (0..head_dim).map(move |dh| qkv_s[t * 3 * d + d + hi * head_dim + dh])
+                    })
+                    .collect();
+                let v: Vec<f32> = (0..n)
+                    .flat_map(|t| {
+                        (0..head_dim).map(move |dh| qkv_s[t * 3 * d + 2 * d + hi * head_dim + dh])
+                    })
+                    .collect();
+                let head_out = sdp_attention_cpu(&q, &k, &v, n, n, head_dim, head_dim);
+                for t in 0..n {
+                    for dh in 0..head_dim {
+                        attn_out[t * d + hi * head_dim + dh] += head_out[t * head_dim + dh];
+                    }
+                }
+            }
+            let proj = linear_cpu(&attn_out, &blk.out_w, &blk.out_b, n, d, d);
+            for i in 0..hidden.len() {
+                hidden[i] += proj[i];
+            }
+
+            let mut normed2 = hidden.clone();
+            layer_norm_rows_cpu(&mut normed2, d);
+            let mut ffn1 = linear_cpu(&normed2, &blk.ffn1_w, &blk.ffn1_b, n, d, ffn);
+            for v in &mut ffn1 {
+                *v = gelu_cpu(*v);
+            }
+            let ffn2 = linear_cpu(&ffn1, &blk.ffn2_w, &blk.ffn2_b, n, ffn, d);
+            for i in 0..hidden.len() {
+                hidden[i] += ffn2[i];
+            }
+        }
+
+        hidden
+    }
+
+    // ── 34. Numerical parity: CPU reference vs GpuViTEncoder (cuda feature) ──
     //
     // With a live CUDA device (`cargo test --features cuda`), GpuViTEncoder
     // runs the whole ViT block GPU-resident (GEMM + layer norm + softmax +
-    // GELU + bias/residual-add + QKV head extract/scatter all as GPU ops).
+    // GELU + bias/residual-add + QKV head extract/scatter all as GPU ops),
-    // This asserts its output matches CpuViTEncoder's reference implementation
+    // over *all* configured blocks (Batch 33: `active_blocks` depth cap
-    // within a tight relative tolerance, for both a single- and multi-token
+    // removed). The reference is `cpu_reference_full_depth` rather than
-    // input and across two different configs.
+    // `CpuViTEncoder::encode` directly: `CpuViTEncoder::encode` (jepa_vit.rs
+    // — a sibling module out of scope for this file) still has its own,
+    // independent `min(depth, 2)` bring-up cap, so comparing against it
+    // directly would only be a valid reference for depth <= 2 configs.
+    // `cpu_reference_full_depth` re-derives the identical per-block math over
+    // every block, so this is a true full-depth comparison for any depth.
     #[cfg(feature = "cuda")]
     #[test]
     fn test_gpu_encoder_matches_cpu_encoder_numerically() {
@@ -1648,7 +1745,7 @@ mod tests {
                 "test requires a live CUDA device with successful weight upload"
             );
 
-            let cpu_out = cpu_enc.encode(patch_indices);
+            let cpu_out = cpu_reference_full_depth(&cpu_enc.cpu_encoder, patch_indices);
             let gpu_out = gpu_enc.encode(patch_indices);
             assert_eq!(cpu_out.len(), gpu_out.len());
 
@@ -1676,4 +1773,57 @@ mod tests {
         check(mini_cfg(), &[0]);
         check(tiny_cfg(), &[0, 5, 10, 15]);
     }
+
+    // ── 35. Full-depth numerical parity (depth=12, embed_dim=192) ────────────
+    //
+    // `GpuViTEncoder::try_encode_gpu` now runs *every* configured block
+    // GPU-resident (the old `active_blocks = blocks.len().min(2)` bring-up
+    // cap has been removed). Since `CpuViTEncoder::encode` itself still has
+    // an independent `min(depth, 2)` cap (in jepa_vit.rs — out of scope for
+    // this file), the reference here is `cpu_reference_full_depth`, which
+    // re-derives the exact same per-block math over all 12 blocks directly
+    // from the block weight arrays, so this is a true full-depth GPU vs CPU
+    // comparison.
+    #[cfg(feature = "cuda")]
+    #[test]
+    fn test_gpu_encoder_matches_cpu_full_depth() {
+        let cfg = tiny_cfg(); // embed_dim=192, depth=12, num_heads=3
+        let gpu_enc = GpuViTEncoder::cuda(cfg.clone(), 0);
+        assert!(
+            gpu_enc.has_gpu_weights(),
+            "test requires a live CUDA device with successful weight upload"
+        );
+        assert_eq!(
+            gpu_enc.gpu_buffer_count(),
+            3 + 8 * cfg.depth,
+            "all {} blocks must have GPU weight buffers uploaded",
+            cfg.depth
+        );
+
+        let patch_indices = [0usize, 5, 10, 15, 20];
+        let cpu_ref = cpu_reference_full_depth(&gpu_enc.cpu_encoder, &patch_indices);
+        let gpu_out = gpu_enc.encode(&patch_indices);
+        assert_eq!(cpu_ref.len(), gpu_out.len());
+        assert_eq!(cpu_ref.len(), patch_indices.len() * cfg.embed_dim);
+
+        let mut max_abs_err = 0.0f32;
+        let mut max_rel_err = 0.0f32;
+        for (c, g) in cpu_ref.iter().zip(gpu_out.iter()) {
+            assert!(c.is_finite() && g.is_finite(), "cpu={c} gpu={g} must be finite");
+            let abs_err = (c - g).abs();
+            let rel_err = abs_err / c.abs().max(1.0);
+            max_abs_err = max_abs_err.max(abs_err);
+            max_rel_err = max_rel_err.max(rel_err);
+        }
+        eprintln!(
+            "[parity/full-depth] depth={} embed_dim={} n={} max_abs_err={max_abs_err:.6} max_rel_err={max_rel_err:.6}",
+            cfg.depth,
+            cfg.embed_dim,
+            patch_indices.len(),
+        );
+        assert!(
+            max_rel_err < 1e-3,
+            "full-depth GPU/CPU ViT diverge: max_rel_err={max_rel_err}, max_abs_err={max_abs_err}"
+        );
+    }
 }

crates/training/rtx-transformers/src/ssl/jepa_runner.rs

@@ -57,6 +57,8 @@ impl RealDataState {
 /// Human-readable ViT size string that parses to [`JepaViTConfig`].
 #[derive(Debug, Clone, PartialEq)]
 pub enum ViTSizeStr {
+    /// Non-standard test/smoke size (d=32, depth=2) — see JepaViTConfig::micro
+    Micro,
     Tiny,
     Small,
     Base,
@@ -68,6 +70,7 @@ impl ViTSizeStr {
     /// Convert to a `JepaViTConfig` with the given image and patch sizes.
     pub fn to_vit_config(&self, image_size: usize, patch_size: usize) -> JepaViTConfig {
         let base = match self {
+            ViTSizeStr::Micro => JepaViTConfig::micro(),
             ViTSizeStr::Tiny  => JepaViTConfig::tiny(),
             ViTSizeStr::Small => JepaViTConfig::small(),
             ViTSizeStr::Base  => JepaViTConfig::base(),
@@ -81,6 +84,7 @@ impl ViTSizeStr {
     /// layout via [`JepaParallelConfig::for_model_and_cluster`].
     pub fn approx_params_m(&self) -> usize {
         match self {
+            ViTSizeStr::Micro => 1,
             ViTSizeStr::Tiny => 6,
             ViTSizeStr::Small => 22,
             ViTSizeStr::Base => 86,
@@ -451,12 +455,13 @@ fn parse_string_array_value(raw: &str, line_no: usize) -> Result<Vec<String>, St
 
 fn parse_vit_size_str(s: &str, line_no: usize) -> Result<ViTSizeStr, String> {
     match s.to_lowercase().as_str() {
+        "micro" => Ok(ViTSizeStr::Micro),
         "tiny"  => Ok(ViTSizeStr::Tiny),
         "small" => Ok(ViTSizeStr::Small),
         "base"  => Ok(ViTSizeStr::Base),
         "large" => Ok(ViTSizeStr::Large),
         "huge"  => Ok(ViTSizeStr::Huge),
-        other   => Err(format!("line {line_no}: unknown vit_size '{other}'; expected tiny|small|base|large|huge")),
+        other   => Err(format!("line {line_no}: unknown vit_size '{other}'; expected micro|tiny|small|base|large|huge")),
     }
 }
 
@@ -531,6 +536,7 @@ impl JepaCheckpoint {
 
         let cfg = &self.config;
         let vit_size_str = match cfg.vit_size {
+            ViTSizeStr::Micro => "micro",
             ViTSizeStr::Tiny  => "tiny",
             ViTSizeStr::Small => "small",
             ViTSizeStr::Base  => "base",
@@ -740,9 +746,17 @@ fn build_encoder(config: &JepaRunConfig, vit_cfg: super::jepa_vit::JepaViTConfig
 ///
 /// Returns aggregate [`JepaTrainingSummary`] after all steps complete.
 pub fn run_jepa_training(config: JepaRunConfig) -> JepaTrainingSummary {
-    // 0. Set up distributed gradient sync
+    // 0. Set up distributed gradient sync. world_size > 1 uses the real
+    // TCP AllReduce backend (rendezvous at master_addr:master_port);
+    // constructor errors are config errors and fail fast.
     let grad_sync = if config.world_size > 1 {
-        JepaGradSync::simulated(config.world_size, config.rank)
+        JepaGradSync::tcp(
+            config.world_size,
+            config.rank,
+            &config.master_addr,
+            config.master_port,
+        )
+        .expect("invalid TCP grad-sync config (world_size/rank/master_port)")
     } else {
         JepaGradSync::single_process()
     };
@@ -863,9 +877,13 @@ pub fn run_jepa_training(config: JepaRunConfig) -> JepaTrainingSummary {
             wall_ms,
         };
 
-        // Simulate gradient sync: scale loss by world_size (effective batch signal)
+        // Gradient sync: real AllReduce over the TCP backend when
+        // world_size > 1 (no-op single-process otherwise). A collective
+        // failure means ranks would diverge, so fail hard.
         let mut fake_grads = vec![step_result.loss; 1]; // representative scalar
-        let _sync_result: GradSyncResult = grad_sync.sync_gradients(&mut fake_grads, step_result.loss);
+        let _sync_result: GradSyncResult = grad_sync
+            .all_reduce(&mut fake_grads, step_result.loss)
+            .expect("gradient AllReduce failed — aborting to avoid rank divergence");
         let _effective_loss = fake_grads[0]; // averaged across world_size
 
         loss_history.push(step_result.loss);
@@ -1010,6 +1028,7 @@ pub fn run_jepa_benchmark(config: &JepaRunConfig) -> BenchmarkResult {
     let config_label = format!(
         "ViT-{}/{} ({})",
         match config.vit_size {
+            ViTSizeStr::Micro => "Micro",
             ViTSizeStr::Tiny  => "Tiny",
             ViTSizeStr::Small => "Small",
             ViTSizeStr::Base  => "Base",
@@ -1075,6 +1094,17 @@ pub fn run_jepa_benchmark(config: &JepaRunConfig) -> BenchmarkResult {
 mod tests {
     use super::*;
 
+    /// Cheap base config for tests that actually run training/benchmark
+    /// steps: ViT-Micro and a tiny batch. Explicit fields in the struct
+    /// literal always override these.
+    fn micro_cfg() -> JepaRunConfig {
+        JepaRunConfig {
+            vit_size: ViTSizeStr::Micro,
+            batch_size: 2,
+            ..JepaRunConfig::default()
+        }
+    }
+
     // 1. Default config sanity
     #[test]
     fn test_default_config_valid() {
@@ -1136,7 +1166,7 @@ mod tests {
             total_steps: 1000,
             warmup_steps: 100,
             base_lr: 1.5e-4,
-            ..JepaRunConfig::default()
+            ..micro_cfg()
         };
         let lr = compute_lr(100, &cfg);
         // At step == warmup_steps: warmup_factor = 1.0, decay_progress = 0 → cosine = 1
@@ -1150,7 +1180,7 @@ mod tests {
             total_steps: 1000,
             warmup_steps: 100,
             base_lr: 1.5e-4,
-            ..JepaRunConfig::default()
+            ..micro_cfg()
         };
         let lr = compute_lr(0, &cfg);
         // warmup_factor = 0/100 = 0 → lr = 0
@@ -1164,7 +1194,7 @@ mod tests {
             total_steps: 1000,
             warmup_steps: 100,
             base_lr: 1.5e-4,
-            ..JepaRunConfig::default()
+            ..micro_cfg()
         };
         let lr = compute_lr(1000, &cfg);
         // decay_progress = 1.0 → cos(π) = -1 → cosine_factor = 0
@@ -1175,13 +1205,13 @@ mod tests {
     #[test]
     fn test_run_tiny_10_steps() {
         let config = JepaRunConfig {
-            vit_size: ViTSizeStr::Tiny,
+            vit_size: ViTSizeStr::Micro,
             total_steps: 10,
             warmup_steps: 2,
             log_every: 100,       // suppress output
             checkpoint_every: 10_000, // no checkpoint during test
             batch_size: 1,        // keep it fast
-            ..JepaRunConfig::default()
+            ..micro_cfg()
         };
         let summary = run_jepa_training(config);
         assert_eq!(summary.total_steps, 10);
@@ -1191,13 +1221,13 @@ mod tests {
     #[test]
     fn test_run_returns_finite_loss() {
         let config = JepaRunConfig {
-            vit_size: ViTSizeStr::Tiny,
+            vit_size: ViTSizeStr::Micro,
             total_steps: 5,
             warmup_steps: 1,
             log_every: 100,
             checkpoint_every: 10_000,
             batch_size: 1,
-            ..JepaRunConfig::default()
+            ..micro_cfg()
         };
         let summary = run_jepa_training(config);
         assert!(summary.final_loss.is_finite(), "final_loss must be finite");
@@ -1208,13 +1238,13 @@ mod tests {
     #[test]
     fn test_run_throughput_positive() {
         let config = JepaRunConfig {
-            vit_size: ViTSizeStr::Tiny,
+            vit_size: ViTSizeStr::Micro,
             total_steps: 3,
             warmup_steps: 1,
             log_every: 100,
             checkpoint_every: 10_000,
             batch_size: 1,
-            ..JepaRunConfig::default()
+            ..micro_cfg()
         };
         let summary = run_jepa_training(config);
         assert!(summary.steps_per_second > 0.0, "steps_per_second must be positive");
@@ -1268,14 +1298,14 @@ mod tests {
     #[test]
     fn test_run_metrics_logger_integration() {
         let config = JepaRunConfig {
-            vit_size: ViTSizeStr::Tiny,
+            vit_size: ViTSizeStr::Micro,
             total_steps: 5,
             warmup_steps: 1,
             log_every: 100,
             checkpoint_every: 10_000,
             batch_size: 1,
             metrics_csv_path: None,
-            ..JepaRunConfig::default()
+            ..micro_cfg()
         };
         let summary = run_jepa_training(config);
         assert!(summary.final_loss.is_finite(), "final_loss must be finite after 5 steps");
@@ -1304,7 +1334,7 @@ mod tests {
             base_lr: 1.5e-4,
             ema_tau_start: 0.996,
             ema_tau_end: 1.0,
-            ..JepaRunConfig::default()
+            ..micro_cfg()
         };
 
         let result = JepaStepResult {
@@ -1348,10 +1378,10 @@ mod tests {
     #[test]
     fn test_benchmark_runs() {
         let config = JepaRunConfig {
-            vit_size: ViTSizeStr::Tiny,
+            vit_size: ViTSizeStr::Micro,
             benchmark_steps: 5,
             benchmark_patches: 4,
-            ..JepaRunConfig::default()
+            ..micro_cfg()
         };
         let result = run_jepa_benchmark(&config);
         assert_eq!(result.steps, 5);
@@ -1362,10 +1392,10 @@ mod tests {
     #[test]
     fn test_benchmark_throughput_positive() {
         let config = JepaRunConfig {
-            vit_size: ViTSizeStr::Tiny,
+            vit_size: ViTSizeStr::Micro,
             benchmark_steps: 3,
             benchmark_patches: 4,
-            ..JepaRunConfig::default()
+            ..micro_cfg()
         };
         let result = run_jepa_benchmark(&config);
         assert!(result.mean_patches_per_sec > 0.0, "throughput must be positive");
@@ -1375,10 +1405,10 @@ mod tests {
     #[test]
     fn test_benchmark_p99_gte_p50() {
         let config = JepaRunConfig {
-            vit_size: ViTSizeStr::Tiny,
+            vit_size: ViTSizeStr::Micro,
             benchmark_steps: 10,
             benchmark_patches: 4,
-            ..JepaRunConfig::default()
+            ..micro_cfg()
         };
         let result = run_jepa_benchmark(&config);
         assert!(
@@ -1397,7 +1427,7 @@ mod tests {
             total_steps: 3,
             log_every: 100,
             checkpoint_every: 10_000,
-            ..JepaRunConfig::default()
+            ..micro_cfg()
         };
         let summary = run_jepa_training(config);
         assert_eq!(summary.total_steps, 3);
@@ -1412,7 +1442,7 @@ mod tests {
             batch_size: 8,
             log_every: 100,
             checkpoint_every: 10_000,
-            ..JepaRunConfig::default()
+            ..micro_cfg()
         };
         let summary = run_jepa_training(config);
         assert!(summary.tokens_per_sec >= 0.0, "tokens_per_sec must be non-negative");
@@ -1425,7 +1455,7 @@ mod tests {
             vit_size: ViTSizeStr::Tiny,
             benchmark_steps: 3,
             benchmark_patches: 4,
-            ..JepaRunConfig::default()
+            ..micro_cfg()
         };
         let r = run_jepa_benchmark(&config);
         assert!(r.config_label.contains("Tiny"), "label '{}' must contain Tiny", r.config_label);
@@ -1439,7 +1469,7 @@ mod tests {
             patch_size: 14,
             benchmark_steps: 3,
             benchmark_patches: 4,
-            ..JepaRunConfig::default()
+            ..micro_cfg()
         };
         let r = run_jepa_benchmark(&config);
         assert!(r.config_label.contains("14"), "label '{}' must contain patch size", r.config_label);
@@ -1455,7 +1485,7 @@ mod tests {
             total_steps: 3,
             log_every: 100,
             checkpoint_every: 10_000,
-            ..JepaRunConfig::default()
+            ..micro_cfg()
         };
         let summary = run_jepa_training(config);
         assert_eq!(summary.total_steps, 3);
@@ -1470,7 +1500,7 @@ mod tests {
             total_steps: 3,
             log_every: 100,
             checkpoint_every: 10_000,
-            ..JepaRunConfig::default()
+            ..micro_cfg()
         };
         // Must not panic — GpuViTEncoder falls back to CPU when no device
         let summary = run_jepa_training(config);
@@ -1485,7 +1515,7 @@ mod tests {
             use_gpu: false,
             benchmark_steps: 3,
             benchmark_patches: 4,
-            ..JepaRunConfig::default()
+            ..micro_cfg()
         };
         let r = run_jepa_benchmark(&config);
         assert!(r.config_label.contains("CPU"), "label '{}' must contain CPU", r.config_label);
@@ -1498,7 +1528,7 @@ mod tests {
             use_gpu: true,
             benchmark_steps: 3,
             benchmark_patches: 4,
-            ..JepaRunConfig::default()
+            ..micro_cfg()
         };
         let r = run_jepa_benchmark(&config);
         assert!(r.mean_patches_per_sec > 0.0);
@@ -1532,7 +1562,7 @@ mod tests {
             total_steps: 3,
             log_every: 100,
             checkpoint_every: 10_000,
-            ..JepaRunConfig::default()
+            ..micro_cfg()
         };
         let summary = run_jepa_training(config);
         assert_eq!(summary.real_data_steps, 0);
@@ -1546,7 +1576,7 @@ mod tests {
             total_steps: 3,
             log_every: 100,
             checkpoint_every: 10_000,
-            ..JepaRunConfig::default()
+            ..micro_cfg()
         };
         let summary = run_jepa_training(config);
         // No shard exists on disk, so real_data_steps=0 (synthetic fallback)
@@ -1561,7 +1591,7 @@ mod tests {
             total_steps: 2,
             log_every: 100,
             checkpoint_every: 10_000,
-            ..JepaRunConfig::default()
+            ..micro_cfg()
         };
         let summary = run_jepa_training(config);
         let _ = summary.real_data_steps; // field must exist
@@ -1575,7 +1605,7 @@ mod tests {
             total_steps: 3,
             log_every: 100,
             checkpoint_every: 10_000,
-            ..JepaRunConfig::default()
+            ..micro_cfg()
         };
         // HTTP shards can't be loaded without network; should fall back to synthetic gracefully
         let summary = run_jepa_training(config);
@@ -1590,7 +1620,7 @@ mod tests {
             total_steps: 3,
             log_every: 100,
             checkpoint_every: 10_000,
-            ..JepaRunConfig::default()
+            ..micro_cfg()
         };
         let summary = run_jepa_training(config);
         assert!(summary.tokens_per_sec >= 0.0);
@@ -1607,7 +1637,7 @@ mod tests {
             total_steps: 2,
             log_every: 100,
             checkpoint_every: 10_000,
-            ..JepaRunConfig::default()
+            ..micro_cfg()
         };
         let summary = run_jepa_training(config);
         assert_eq!(summary.total_steps, 2);

crates/training/rtx-transformers/src/ssl/jepa_vit.rs

@@ -71,6 +71,13 @@ pub struct JepaViTConfig {
 
 impl JepaViTConfig {
     /// ViT-Tiny: d=192, depth=12, heads=3, patch=16, img=224
+    /// ViT-Micro: d=32, depth=2, heads=4, patch=16, img=64 — not a standard
+    /// size; intended for tests and smoke runs where full-depth Tiny is
+    /// needlessly slow on CPU.
+    pub fn micro() -> Self {
+        Self { embed_dim: 32, depth: 2, num_heads: 4, mlp_ratio: 2.0, patch_size: 16, image_size: 64 }
+    }
+
     pub fn tiny() -> Self {
         Self { embed_dim: 192, depth: 12, num_heads: 3, mlp_ratio: 4.0, patch_size: 16, image_size: 224 }
     }
@@ -340,8 +347,8 @@ impl ViTBlock {
 ///
 /// Maintains a flat patch embedding table initialised with sinusoidal
 /// position encodings plus per-patch linear projection.  Each `encode`
-/// call runs `min(depth, 2)` ViT blocks to keep CPU tests fast while
+/// call runs all `depth` ViT blocks (tests wanting speed should use a
-/// still exercising the full architectural pathway.
+/// small-depth config rather than relying on a hidden cap).
 #[derive(Debug, Clone)]
 pub struct CpuViTEncoder {
     pub config: JepaViTConfig,
@@ -350,7 +357,7 @@ pub struct CpuViTEncoder {
     /// Patch projection weights [embed_dim, embed_dim]
     pub(crate) proj_w: Vec<f32>,
     pub(crate) proj_b: Vec<f32>,
-    /// Transformer blocks (full depth stored, but only min(depth,2) run)
+    /// Transformer blocks (full depth)
     pub(crate) blocks: Vec<ViTBlock>,
 }
 
@@ -475,10 +482,9 @@ impl JepaEncoder for CpuViTEncoder {
             }
         }
 
-        // Run min(depth, 2) ViT blocks
+        // Run all ViT blocks (full depth)
-        let active_blocks = self.blocks.len().min(2);
         let mut hidden = projected;
-        for blk in self.blocks.iter().take(active_blocks) {
+        for blk in &self.blocks {
             hidden = blk.forward(&hidden, n);
         }
         hidden