feat(jepa): eval in the training loop, NCCL GPU AllReduce, GPU checkpointing

- Eval: run_jepa_training now runs k-NN (k=5) + linear-probe evaluation
  every eval_every steps and at the end (deduped when aligned);
  JepaEvalResult recorded in JepaTrainingSummary (final_knn_acc /
  final_probe_acc), printed by the CLI, appended as an # eval section
  to the metrics CSV. Probe set is deterministic LCG synthetic (offset
  seed, never aliases training batches) or real shard labels when
  loaded.
- NCCL: real GPU-direct AllReduce backend behind the new `nccl`
  feature (cudarc/nccl, dlopen-based so builds don't need libnccl).
  NCCL unique id is bootstrapped over the existing TCP rendezvous
  (master_port+137); data path is htod -> ncclAllReduce(Sum) -> dtoh
  -> mean. catch_unwind guards cudarc's panic-on-missing-lib so
  training falls back instead of aborting. Verified for real on the
  RTX 5060 Ti: single-rank GPU all_reduce identity test passes
  (26/26 with --features nccl).
- GPU checkpointing/eval: JepaTrainerV2::context_encoder_cpu_weights()
  exposes host-side weights for both CPU and GPU encoders
  (GpuViTEncoder::cpu_weights); checkpoint save and eval now work for
  GPU training runs (verified: .jepa binaries written and 2 eval
  passes during a live GPU CLI run). Resume with a GPU encoder
  restores the step counter and warns that weight re-upload is not
  yet implemented rather than silently training on stale weights.

125 runner/distributed/vit tests pass; CLI 8/8; cuda check clean.

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

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

@@ -192,6 +192,11 @@ fn print_training_summary(summary: &rtx_transformers::ssl::JepaTrainingSummary)
     println!("World size:           {}", summary.world_size);
     println!("Effective batch:      {}", summary.effective_batch_size);
     println!("Real-data steps:      {}", summary.real_data_steps);
+    println!("Eval passes:          {}", summary.eval_results.len());
+    if !summary.eval_results.is_empty() {
+        println!("Final knn_acc:        {:.4}", summary.final_knn_acc);
+        println!("Final probe_acc:      {:.4}", summary.final_probe_acc);
+    }
 }
 
 fn cmd_bench(vit: VitSizeArg, steps: usize, gpu: bool, metrics_csv: Option<String>) -> Result<()> {

crates/training/rtx-transformers/Cargo.toml

@@ -87,6 +87,12 @@ default = []
 cuda = ["cudarc", "rtx-flash-attention/cuda", "rtx-tensor/cuda", "rtx-runtime/cuda"]
 metal = ["rtx-flash-attention/metal", "rtx-tensor/metal", "rtx-runtime/metal", "dep:objc2", "dep:objc2-metal", "dep:objc2-foundation", "dep:block2"]
 cpu = ["rtx-tensor/cpu"]
+# Real NCCL GPU-direct AllReduce for JEPA gradient sync (jepa_distributed.rs).
+# Requires a CUDA GPU and libnccl.so discoverable at runtime (cudarc uses
+# dynamic-loading/dlopen for NCCL by default, so this feature builds fine
+# without libnccl present; it only fails at the point a `GradSyncBackend::Nccl`
+# handle is actually used at runtime).
+nccl = ["cuda", "cudarc/nccl"]
 disabled_tests = []
 vision-bridge = ["rtx-vision"]
 # Real JPEG/PNG pixel decoding for WebDataset records; without it,

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

@@ -312,6 +312,91 @@ mod tcp_allreduce {
         Ok(())
     }
 
+    /// Broadcast an arbitrary byte payload from rank 0 to every other rank,
+    /// using the same rank/world_size handshake as [`all_reduce`]. This is
+    /// the standard out-of-band NCCL bootstrap pattern: rank 0 creates the
+    /// 128-byte NCCL unique id and every other rank must learn it before
+    /// calling `Comm::from_rank`, since NCCL itself has no notion of a
+    /// network rendezvous — only this out-of-band channel does.
+    ///
+    /// `root_payload` must be `Some` on rank 0 (the value to broadcast) and
+    /// is ignored elsewhere. Returns the broadcast payload on every rank.
+    #[cfg(feature = "nccl")]
+    pub(super) fn broadcast_from_root(
+        world_size: usize,
+        rank: usize,
+        master_addr: &str,
+        master_port: u16,
+        root_payload: Option<Vec<u8>>,
+        timeout_secs: u64,
+    ) -> Result<Vec<u8>, String> {
+        let timeout = Duration::from_secs(timeout_secs);
+        if rank == 0 {
+            let payload = root_payload
+                .ok_or_else(|| "rank 0 must supply root_payload for broadcast".to_string())?;
+            let bind_addr = format!("{master_addr}:{master_port}");
+            let listener = TcpListener::bind(&bind_addr)
+                .map_err(|e| format!("rank 0 bootstrap bind failed on {bind_addr}: {e}"))?;
+
+            let mut seen_ranks: std::collections::HashSet<u32> = std::collections::HashSet::new();
+            let mut accepted = 0usize;
+            while accepted < 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();
+
+                let peer_rank = read_u32(&mut stream).map_err(|e| {
+                    format!("bootstrap handshake read rank from {peer_addr} failed: {e}")
+                })?;
+                let peer_world_size = read_u32(&mut stream).map_err(|e| {
+                    format!("bootstrap handshake read world_size from {peer_addr} failed: {e}")
+                })?;
+                if peer_world_size as usize != world_size
+                    || peer_rank == 0
+                    || peer_rank as usize >= world_size
+                    || !seen_ranks.insert(peer_rank)
+                {
+                    return Err(format!(
+                        "bootstrap handshake rejected peer {peer_addr} (rank {peer_rank}, world_size {peer_world_size})"
+                    ));
+                }
+
+                write_len_prefixed(&mut stream, &payload)
+                    .map_err(|e| format!("bootstrap send to {peer_addr} failed: {e}"))?;
+                accepted += 1;
+            }
+            Ok(payload)
+        } else {
+            let addr = format!("{master_addr}:{master_port}");
+            let sock_addr: SocketAddr = addr
+                .parse()
+                .map_err(|e| format!("parse master addr '{addr}': {e}"))?;
+            let deadline = Instant::now() + timeout;
+            let mut stream = loop {
+                match TcpStream::connect_timeout(&sock_addr, timeout) {
+                    Ok(s) => break s,
+                    Err(e) => {
+                        if Instant::now() >= deadline {
+                            return Err(format!(
+                                "rank {rank} bootstrap 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();
+
+            write_u32(&mut stream, rank as u32)
+                .map_err(|e| format!("bootstrap handshake send rank failed: {e}"))?;
+            write_u32(&mut stream, world_size as u32)
+                .map_err(|e| format!("bootstrap handshake send world_size failed: {e}"))?;
+
+            read_len_prefixed(&mut stream).map_err(|e| format!("bootstrap recv failed: {e}"))
+        }
+    }
+
     /// Entry point: dispatches to `run_primary` or `run_peer` based on rank.
     pub(super) fn all_reduce(
         world_size: usize,
@@ -461,6 +546,126 @@ impl TcpRingAllReduce {
     }
 }
 
+// ============================================================================
+// nccl_gpu — real GPU-direct NCCL AllReduce (feature-gated)
+// ============================================================================
+//
+// Bootstrap: NCCL has no notion of network rendezvous by itself — every rank
+// must be handed the same 128-byte "unique id" before calling
+// `Comm::from_rank`. The standard pattern (used by PyTorch's
+// `c10d::ProcessGroupNCCL`, Horovod, etc.) is: rank 0 creates the id and
+// broadcasts it out-of-band. We reuse the `tcp_allreduce` rendezvous
+// machinery in this same file for that broadcast (on a dedicated port offset
+// from the gradient-exchange port so the two rounds never collide), then
+// hand the id to `cudarc::nccl::Comm::from_rank` on every rank.
+//
+// Data path: upload the local f32 gradient vector to a `CudaSlice` on the
+// rank's GPU, run a real `ncclAllReduce` (Sum) between GPU-resident buffers
+// (GPU-direct between comm members — no host round-trip mid-collective),
+// download the summed result, then divide by `world_size` on the host. The
+// host copies at the upload/download edges are intentional for now (see
+// module docs); the actual communication is real NCCL, not simulated.
+#[cfg(feature = "nccl")]
+mod nccl_gpu {
+    use super::tcp_allreduce;
+    use cudarc::driver::CudaContext;
+    use cudarc::nccl::{Comm, Id, ReduceOp};
+
+    /// Timeout for the id-bootstrap rendezvous; short in tests so a missing
+    /// peer fails fast instead of hanging the suite.
+    #[cfg(test)]
+    const BOOTSTRAP_TIMEOUT_SECS: u64 = 2;
+    #[cfg(not(test))]
+    const BOOTSTRAP_TIMEOUT_SECS: u64 = 30;
+
+    /// Offset applied to `master_port` for the id-bootstrap rendezvous so it
+    /// never collides with the gradient-exchange port used immediately after.
+    const ID_BOOTSTRAP_PORT_OFFSET: u16 = 137;
+
+    /// Perform a real GPU-resident NCCL AllReduce (sum) on `grads`, then
+    /// divide by `world_size` in place. Returns `Err` on any bootstrap,
+    /// CUDA, or NCCL failure so the caller can fall back to simulated
+    /// averaging instead of hard-blocking training.
+    pub(super) fn gpu_all_reduce_average(
+        world_size: usize,
+        rank: usize,
+        master_addr: &str,
+        master_port: u16,
+        grads: &mut Vec<f32>,
+    ) -> Result<(), String> {
+        // 1. Bootstrap the NCCL unique id out-of-band over TCP (rank 0
+        //    creates it, every rank learns it via the same rendezvous
+        //    machinery `tcp_allreduce` uses for gradient exchange).
+        let bootstrap_port = master_port
+            .checked_add(ID_BOOTSTRAP_PORT_OFFSET)
+            .ok_or_else(|| "master_port + bootstrap offset overflowed u16".to_string())?;
+
+        let id_bytes = if rank == 0 {
+            let id = Id::new().map_err(|e| format!("Id::new failed: {e:?}"))?;
+            let raw: Vec<u8> = id.internal().iter().map(|&c| c as u8).collect();
+            tcp_allreduce::broadcast_from_root(
+                world_size,
+                rank,
+                master_addr,
+                bootstrap_port,
+                Some(raw.clone()),
+                BOOTSTRAP_TIMEOUT_SECS,
+            )?;
+            raw
+        } else {
+            tcp_allreduce::broadcast_from_root(
+                world_size,
+                rank,
+                master_addr,
+                bootstrap_port,
+                None,
+                BOOTSTRAP_TIMEOUT_SECS,
+            )?
+        };
+
+        if id_bytes.len() != 128 {
+            return Err(format!(
+                "expected a 128-byte NCCL unique id, got {} bytes",
+                id_bytes.len()
+            ));
+        }
+        let mut internal = [0 as std::ffi::c_char; 128];
+        for (dst, &b) in internal.iter_mut().zip(id_bytes.iter()) {
+            *dst = b as std::ffi::c_char;
+        }
+        let id = Id::uninit(internal);
+
+        // 2. Initialise the CUDA context + NCCL communicator for this rank.
+        //    Single-GPU-per-node (this environment): ordinal 0 for every
+        //    rank. On a multi-GPU node this would be `rank % local_gpus`.
+        let ctx =
+            CudaContext::new(0).map_err(|e| format!("CudaContext::new(0) failed: {e:?}"))?;
+        let stream = ctx.default_stream();
+        let comm = Comm::from_rank(stream.clone(), rank, world_size, id)
+            .map_err(|e| format!("Comm::from_rank failed: {e:?}"))?;
+
+        // 3. Upload -> ncclAllReduce(Sum) -> download -> divide by world_size.
+        let send_buf = stream
+            .clone_htod(grads.as_slice())
+            .map_err(|e| format!("upload gradients to GPU failed: {e:?}"))?;
+        let mut recv_buf = unsafe { stream.alloc::<f32>(grads.len()) }
+            .map_err(|e| format!("GPU recv buffer alloc failed: {e:?}"))?;
+
+        comm.all_reduce(&send_buf, &mut recv_buf, &ReduceOp::Sum)
+            .map_err(|e| format!("NCCL all_reduce failed: {e:?}"))?;
+
+        let mut host_out: Vec<f32> = stream
+            .clone_dtoh(&recv_buf)
+            .map_err(|e| format!("download reduced gradients from GPU failed: {e:?}"))?;
+
+        for v in &mut host_out {
+            *v /= world_size as f32;
+        }
+        *grads = host_out;
+        Ok(())
+    }
+}
+
 // ============================================================================
 // GradSyncBackend
 // ============================================================================
@@ -470,9 +675,20 @@ impl TcpRingAllReduce {
 pub enum GradSyncBackend {
     /// No actual communication — used for single-process or testing.
     Simulated,
-    /// NCCL/TCP backend: attempts real TCP ring-reduce when the
+    /// GPU-direct NCCL AllReduce backend.
-    /// `distributed-tcp` feature is enabled; falls back to simulated divide
+    ///
-    /// on any network error.
+    /// - With the `nccl` cargo feature: bootstraps a real
+    ///   `cudarc::nccl::Comm` on every rank (rank 0's unique id is
+    ///   broadcast out-of-band over TCP via the `tcp_allreduce` rendezvous
+    ///   helpers in this file) and performs a real GPU-resident
+    ///   `ncclAllReduce(Sum)`, then divides by `world_size`. Falls back to
+    ///   simulated averaging only if the bootstrap/CUDA/NCCL call itself
+    ///   fails at runtime (e.g. no GPU, or `libnccl.so` not found).
+    /// - Without the `nccl` feature: this is a **stub** — it does no
+    ///   communication at all and simply divides by `world_size` in place,
+    ///   identical to `Simulated`. Build with `--features nccl` (which
+    ///   requires a CUDA-capable GPU and `libnccl.so` discoverable at
+    ///   runtime) to get the real behavior.
     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
@@ -526,6 +742,14 @@ impl JepaGradSync {
     }
 
     /// NCCL config (validates world_size > 0 and port > 0).
+    ///
+    /// Construction always succeeds the same way regardless of feature
+    /// flags; what differs is what [`Self::sync_gradients`] actually does
+    /// with it. **Without** the `nccl` cargo feature this is a stub: no
+    /// hardware communication happens and gradients are just divided by
+    /// `world_size` in-process (identical to [`Self::simulated`]). **With**
+    /// `--features nccl` (requires a CUDA GPU and `libnccl.so` discoverable
+    /// at runtime) it performs a real GPU-direct `ncclAllReduce`.
     pub fn nccl(
         world_size: usize,
         rank: usize,
@@ -654,30 +878,82 @@ impl JepaGradSync {
                 }
             }
             GradSyncBackend::Nccl { master_addr, master_port } => {
-                #[cfg(feature = "distributed-tcp")]
+                #[cfg(feature = "nccl")]
                 {
-                    let ring = TcpRingAllReduce {
+                    // cudarc's dynamic-loading NCCL backend *panics* (rather
-                        world_size: self.world_size,
+                    // than returning an `Err`) when `libnccl.so` can't be
-                        rank: self.rank,
+                    // found at runtime, so a `catch_unwind` is required here
-                        master_addr: master_addr.clone(),
+                    // to preserve the "never hard-block training" contract
-                        master_port: *master_port,
+                    // even when the `nccl` feature is compiled in but the
+                    // library isn't installed/discoverable on this host.
+                    let world_size = self.world_size;
+                    let rank = self.rank;
+                    let outcome = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
+                        nccl_gpu::gpu_all_reduce_average(
+                            world_size,
+                            rank,
+                            master_addr,
+                            *master_port,
+                            grads,
+                        )
+                    }));
+                    let failure: Option<String> = match outcome {
+                        Ok(Ok(())) => None,
+                        Ok(Err(e)) => Some(e),
+                        Err(panic_payload) => Some(
+                            panic_payload
+                                .downcast_ref::<String>()
+                                .cloned()
+                                .or_else(|| {
+                                    panic_payload
+                                        .downcast_ref::<&str>()
+                                        .map(|s| (*s).to_string())
+                                })
+                                .unwrap_or_else(|| "unknown panic in NCCL backend".to_string()),
+                        ),
                     };
-                    if let Err(e) = ring.allreduce_sum_divide(grads) {
+                    if let Some(e) = failure {
-                        // Network failure: log warning and fall back to simulated averaging
+                        // Bootstrap/CUDA/NCCL failure (or missing
+                        // libnccl.so): log warning and fall back to
+                        // simulated averaging so the training loop is never
+                        // hard-blocked.
                         eprintln!(
-                            "NCCL TCP allreduce failed (falling back to simulated): {e}"
+                            "NCCL GPU-direct AllReduce failed (falling back to simulated): {e}"
                         );
                         for g in grads.iter_mut() {
-                            *g /= self.world_size as f32;
+                            *g /= world_size as f32;
                         }
                     }
                 }
-                #[cfg(not(feature = "distributed-tcp"))]
+                #[cfg(not(feature = "nccl"))]
                 {
-                    // Without the feature, simulate divide by world_size
+                    #[cfg(feature = "distributed-tcp")]
-                    let _ = (master_addr, master_port); // suppress unused warnings
+                    {
-                    for g in grads.iter_mut() {
+                        let ring = TcpRingAllReduce {
-                        *g /= self.world_size as f32;
+                            world_size: self.world_size,
+                            rank: self.rank,
+                            master_addr: master_addr.clone(),
+                            master_port: *master_port,
+                        };
+                        if let Err(e) = ring.allreduce_sum_divide(grads) {
+                            // Network failure: log warning and fall back to simulated averaging
+                            eprintln!(
+                                "NCCL TCP allreduce failed (falling back to simulated): {e}"
+                            );
+                            for g in grads.iter_mut() {
+                                *g /= self.world_size as f32;
+                            }
+                        }
+                    }
+                    #[cfg(not(feature = "distributed-tcp"))]
+                    {
+                        // Stub (see `GradSyncBackend::Nccl` docs): no
+                        // hardware communication without the `nccl` feature,
+                        // simulate divide by world_size.
+                        let _ = (master_addr, master_port); // suppress unused warnings
+                        for g in grads.iter_mut() {
+                            *g /= self.world_size as f32;
+                        }
                     }
                 }
                 let comm_bytes = (grads.len() * std::mem::size_of::<f32>()) as u64
@@ -1155,3 +1431,38 @@ mod tests {
         );
     }
 }
+
+// ============================================================================
+// Real GPU NCCL tests (only compiled/run with `--features nccl`)
+// ============================================================================
+
+#[cfg(all(test, feature = "nccl"))]
+mod nccl_gpu_tests {
+    use super::nccl_gpu::gpu_all_reduce_average;
+
+    // A world_size=1 NCCL communicator is valid (NCCL supports single-rank
+    // groups) and exercises the real path end to end: CUDA context init,
+    // `Comm::from_rank`, a GPU-resident `ncclAllReduce(Sum)`, and the
+    // host-side divide-by-world_size — with no TCP peer required (rank 0
+    // never waits on `accepted < world_size - 1 == 0`). Sum of one rank
+    // divided by world_size=1 must leave the buffer unchanged.
+    #[test]
+    fn test_nccl_single_rank_real_gpu_allreduce_is_identity() {
+        let mut grads = vec![1.0f32, -2.5, 3.25, 0.0, 42.0];
+        let original = grads.clone();
+
+        let result = gpu_all_reduce_average(1, 0, "127.0.0.1", 31900, &mut grads);
+
+        assert!(
+            result.is_ok(),
+            "single-rank real GPU NCCL all_reduce failed: {:?}",
+            result
+        );
+        for (got, want) in grads.iter().zip(original.iter()) {
+            assert!(
+                (got - want).abs() < 1e-6,
+                "expected {want}, got {got} (world_size=1 all_reduce must be identity)"
+            );
+        }
+    }
+}

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

@@ -441,6 +441,12 @@ pub struct GpuViTEncoder {
 }
 
 impl GpuViTEncoder {
+    /// Host-side copy of the encoder weights (the GPU buffers are uploaded
+    /// from these). Used for checkpointing GPU training runs.
+    pub fn cpu_weights(&self) -> &CpuViTEncoder {
+        &self.cpu_encoder
+    }
+
     /// Create targeting CPU (equivalent to `CpuViTEncoder`, useful for testing).
     pub fn cpu(config: JepaViTConfig) -> Self {
         Self {

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

@@ -580,7 +580,7 @@ impl JepaCheckpoint {
 
         // Binary weight checkpoint (if trainer with CPU encoder is provided)
         if let Some(t) = trainer {
-            if let Some(cpu_enc) = t.context_encoder_as_cpu() {
+            if let Some(cpu_enc) = t.context_encoder_cpu_weights() {
                 let fields = encoder_to_fields(cpu_enc);
                 let config_summary = format!(
                     "{{\"vit_size\":\"{vit_size}\",\"image_size\":{image_size},\"patch_size\":{patch_size},\"total_steps\":{total_steps}}}",
@@ -681,6 +681,162 @@ pub struct JepaTrainingSummary {
     pub tokens_per_sec: f64,
     /// Number of steps that used real (non-synthetic) data
     pub real_data_steps: usize,
+    /// All evaluation passes recorded during the run (see [`JepaEvalResult`])
+    pub eval_results: Vec<JepaEvalResult>,
+    /// k-NN top-1 accuracy from the last evaluation pass (0.0 if no eval ran)
+    pub final_knn_acc: f64,
+    /// Linear-probe top-1 accuracy from the last evaluation pass (0.0 if no eval ran)
+    pub final_probe_acc: f64,
+}
+
+// ============================================================================
+// Evaluation: wired into the training loop every `config.eval_every` steps
+// ============================================================================
+
+/// Number of synthetic classes cycled through when labelling the probe set.
+const EVAL_NUM_CLASSES: usize = 4;
+/// Total size of the small held-out probe set used for k-NN / linear-probe eval.
+/// Kept tiny — this is a training-loop wiring check, not a rigorous benchmark.
+const EVAL_PROBE_SAMPLES: usize = 32;
+/// k for the k-NN evaluation pass.
+const EVAL_KNN_K: usize = 5;
+/// Linear-probe SGD epochs per eval pass.
+const EVAL_PROBE_EPOCHS: usize = 20;
+/// Linear-probe learning rate per eval pass.
+const EVAL_PROBE_LR: f64 = 0.1;
+/// LCG seed for the synthetic probe-set patch selection. Deliberately distinct
+/// from the training synthetic-batch seed (`config.batch_size`, see
+/// `synthetic_batch_size_one`) so the probe set never aliases training data.
+const EVAL_SEED_OFFSET: u64 = 0x9E37_79B9_7F4A_7C15;
+
+/// Result of one evaluation pass during training.
+#[derive(Debug, Clone)]
+pub struct JepaEvalResult {
+    /// Training step at which this evaluation was run
+    pub step: usize,
+    /// k-NN top-1 accuracy on the held-out probe set (0.0–1.0)
+    pub knn_acc: f64,
+    /// Linear-probe top-1 accuracy on the held-out probe set (0.0–1.0)
+    pub probe_acc: f64,
+}
+
+/// Build a small deterministic probe set: `(features, labels)` where
+/// `features` is a flat `[EVAL_PROBE_SAMPLES, embed_dim]` matrix.
+///
+/// Patch indices are chosen via an LCG seeded from [`EVAL_SEED_OFFSET`] — a
+/// distinct seed from the training synthetic-data generator, so this never
+/// reproduces a training sample. Labels come from a freshly pulled batch of
+/// the real data pipeline when one is loaded and carries labels; otherwise
+/// (no real data, or records without labels) labels cycle deterministically
+/// through `EVAL_NUM_CLASSES` synthetic classes.
+fn build_probe_set(
+    encoder: &super::jepa_vit::CpuViTEncoder,
+    real_data: Option<&mut RealDataState>,
+) -> (Vec<f32>, Vec<usize>) {
+    use super::jepa_vit::JepaEncoder as _;
+    let embed_dim = encoder.embed_dim();
+    let num_patches = encoder.num_patches().max(1);
+    let n = EVAL_PROBE_SAMPLES;
+
+    let mut labels: Vec<usize> = Vec::new();
+    if let Some(rd) = real_data {
+        if let Some(batch) = rd.pipeline.next_batch() {
+            if !batch.labels.is_empty() {
+                for i in 0..n {
+                    let lbl = batch.labels[i % batch.labels.len()]
+                        .unwrap_or(i % EVAL_NUM_CLASSES);
+                    labels.push(lbl);
+                }
+            }
+        }
+    }
+    if labels.is_empty() {
+        labels = (0..n).map(|i| i % EVAL_NUM_CLASSES).collect();
+    }
+
+    let mut lcg: u64 = EVAL_SEED_OFFSET;
+    let mut features = Vec::with_capacity(n * embed_dim);
+    for _ in 0..n {
+        lcg = lcg
+            .wrapping_mul(6364136223846793005)
+            .wrapping_add(1442695040888963407);
+        let patch_idx = ((lcg >> 11) as usize) % num_patches;
+        let enc = encoder.encode(&[patch_idx]);
+        features.extend_from_slice(&enc);
+    }
+
+    (features, labels)
+}
+
+/// Run one evaluation pass at `step`: encode the probe set with the current
+/// context encoder, split it in half (first half = k-NN bank / linear-probe
+/// train, second half = query / test), and compute both a k-NN top-1
+/// accuracy and a linear-probe top-1 accuracy.
+///
+/// Returns `None` (evaluation skipped, not a training failure) when the
+/// trainer's context encoder isn't a `CpuViTEncoder` we can introspect (e.g.
+/// an opaque GPU encoder) — evaluation degrades gracefully rather than
+/// aborting the run.
+fn run_eval_pass(
+    trainer: &JepaTrainerV2,
+    real_data: Option<&mut RealDataState>,
+    step: usize,
+) -> Option<JepaEvalResult> {
+    use super::jepa_vit::JepaEncoder as _;
+    let encoder = trainer.context_encoder_cpu_weights()?;
+    let embed_dim = encoder.embed_dim();
+
+    let (features, labels) = build_probe_set(encoder, real_data);
+    let n = labels.len();
+    if n < 4 {
+        return None;
+    }
+
+    let half = n / 2;
+    let (train_feat, test_feat) = features.split_at(half * embed_dim);
+    let (train_labels, test_labels) = labels.split_at(half);
+
+    let evaluator = super::jepa::JepaEvaluator::new(embed_dim);
+    let k = EVAL_KNN_K.min(half.max(1));
+    let knn_res = evaluator.knn_eval(train_feat, train_labels, test_feat, test_labels, k);
+    let probe_res = evaluator.linear_probe(
+        train_feat,
+        train_labels,
+        test_feat,
+        test_labels,
+        EVAL_NUM_CLASSES,
+        EVAL_PROBE_EPOCHS,
+        EVAL_PROBE_LR,
+    );
+
+    Some(JepaEvalResult {
+        step,
+        knn_acc: knn_res.top1_accuracy as f64,
+        probe_acc: probe_res.top1_accuracy as f64,
+    })
+}
+
+/// Append an eval-results section to a metrics CSV file previously written by
+/// `JepaMetricsLogger::save_csv`. Uses a clearly delimited extra section
+/// (blank line + `# eval` marker) rather than merging columns onto the
+/// per-step table, since eval steps are a strict subset of training steps
+/// and most step rows would otherwise carry blank/NaN eval columns.
+fn append_eval_csv_section(csv_path: &str, evals: &[JepaEvalResult]) -> Result<(), String> {
+    use std::io::Write;
+    let mut out = String::new();
+    out.push('\n');
+    out.push_str("# eval\n");
+    out.push_str("step,knn_acc,probe_acc\n");
+    for e in evals {
+        out.push_str(&format!("{},{:.6},{:.6}\n", e.step, e.knn_acc, e.probe_acc));
+    }
+    let mut f = std::fs::OpenOptions::new()
+        .append(true)
+        .create(true)
+        .open(csv_path)
+        .map_err(|e| format!("Cannot open '{csv_path}' for append: {e}"))?;
+    f.write_all(out.as_bytes())
+        .map_err(|e| format!("Cannot write eval section to '{csv_path}': {e}"))
 }
 
 // ============================================================================
@@ -773,8 +929,18 @@ pub fn run_jepa_training(config: JepaRunConfig) -> JepaTrainingSummary {
         match load_checkpoint(resume_path) {
             Ok(ckpt) => {
                 start_step = ckpt.step + 1;
+                // NOTE: weight restore currently only applies to the CPU
+                // encoder. Restoring into a GpuViTEncoder would require
+                // re-uploading the device buffers after mutation (not yet
+                // implemented) — mutating only its host copy would silently
+                // train on stale GPU weights, so we warn instead.
                 if let Some(cpu_enc) = trainer.context_encoder_as_cpu_mut() {
                     let _ = apply_fields_to_encoder(cpu_enc, &ckpt.fields);
+                } else {
+                    eprintln!(
+                        "Warning: resume with a GPU encoder restores the step counter \
+                         but not model weights (GPU weight re-upload not implemented)"
+                    );
                 }
                 eprintln!("Resumed from step {}", ckpt.step);
             }
@@ -837,6 +1003,7 @@ pub fn run_jepa_training(config: JepaRunConfig) -> JepaTrainingSummary {
     let mut loss_history: Vec<f32> = Vec::with_capacity(config.total_steps);
     let mut checkpoints_saved = 0usize;
     let mut real_data_steps = 0usize;
+    let mut eval_results: Vec<JepaEvalResult> = Vec::new();
     let training_start = Instant::now();
     let mut step_start = Instant::now();
 
@@ -928,12 +1095,48 @@ pub fn run_jepa_training(config: JepaRunConfig) -> JepaTrainingSummary {
 
         let _ = step_start; // suppress lint
         step_start = Instant::now();
+
+        // Evaluation pass — only rank 0 evaluates. Runs on the regular
+        // `eval_every` cadence; the post-loop final eval below is skipped
+        // when it would duplicate the cadence eval already run at the last
+        // step (see note there).
+        if grad_sync.is_primary() && config.eval_every > 0 && step % config.eval_every == 0 {
+            if let Some(result) = run_eval_pass(&trainer, real_data.as_mut(), step) {
+                println!(
+                    "[eval] step={} knn_acc={:.4} probe_acc={:.4}",
+                    result.step, result.knn_acc, result.probe_acc
+                );
+                eval_results.push(result);
+            }
+        }
+    }
+
+    // Final evaluation pass after the last step. Design choice: if
+    // `eval_every` divides `total_steps`, the cadence eval above already ran
+    // at `step == total_steps` — running it again here would just duplicate
+    // that exact same probe pass (same encoder state, same deterministic
+    // probe set), so we skip it rather than double-count. Otherwise (e.g.
+    // `total_steps` isn't a multiple of `eval_every`, or `eval_every == 0`),
+    // this is the only evaluation at the final encoder state and always runs.
+    let final_step_already_evaluated =
+        config.eval_every > 0 && config.total_steps % config.eval_every == 0;
+    if grad_sync.is_primary() && config.eval_every > 0 && !final_step_already_evaluated {
+        if let Some(result) = run_eval_pass(&trainer, real_data.as_mut(), config.total_steps) {
+            println!(
+                "[eval] step={} knn_acc={:.4} probe_acc={:.4} (final)",
+                result.step, result.knn_acc, result.probe_acc
+            );
+            eval_results.push(result);
+        }
     }
 
     // Export metrics CSV if requested
     if grad_sync.is_primary() {
         if let Some(ref csv_path) = config.metrics_csv_path {
             let _ = metrics_logger.save_csv(csv_path);
+            if !eval_results.is_empty() {
+                let _ = append_eval_csv_section(csv_path, &eval_results);
+            }
         }
     }
 
@@ -948,6 +1151,11 @@ pub fn run_jepa_training(config: JepaRunConfig) -> JepaTrainingSummary {
         0.0
     };
 
+    let (final_knn_acc, final_probe_acc) = eval_results
+        .last()
+        .map(|e| (e.knn_acc, e.probe_acc))
+        .unwrap_or((0.0, 0.0));
+
     JepaTrainingSummary {
         total_steps: report.total_steps,
         final_loss: report.final_loss,
@@ -959,6 +1167,9 @@ pub fn run_jepa_training(config: JepaRunConfig) -> JepaTrainingSummary {
         effective_batch_size,
         tokens_per_sec,
         real_data_steps,
+        eval_results,
+        final_knn_acc,
+        final_probe_acc,
     }
 }
 
@@ -1642,4 +1853,95 @@ mod tests {
         let summary = run_jepa_training(config);
         assert_eq!(summary.total_steps, 2);
     }
+
+    // ---- Evaluation wiring tests (E1–E4) ----
+
+    // E1. total_steps=4, eval_every=2 → cadence evals at steps 2 and 4; the
+    // post-loop final eval is skipped because eval_every (2) divides
+    // total_steps (4) — see `final_step_already_evaluated` in
+    // `run_jepa_training`. This is the deliberate dedup choice: 2 evals
+    // recorded total, at steps [2, 4], not 3.
+    #[test]
+    fn test_eval_cadence_and_final_dedup() {
+        let config = JepaRunConfig {
+            total_steps: 4,
+            eval_every: 2,
+            log_every: 100,
+            checkpoint_every: 10_000,
+            ..micro_cfg()
+        };
+        let summary = run_jepa_training(config);
+        let steps: Vec<usize> = summary.eval_results.iter().map(|e| e.step).collect();
+        assert_eq!(steps, vec![2, 4], "expected exactly one eval at each cadence step, deduped with the final eval");
+        assert_eq!(summary.eval_results.len(), 2);
+    }
+
+    // E2. When eval_every does not divide total_steps, the final eval is a
+    // distinct extra pass (not deduped).
+    #[test]
+    fn test_eval_final_not_deduped_when_not_multiple() {
+        let config = JepaRunConfig {
+            total_steps: 5,
+            eval_every: 2,
+            log_every: 100,
+            checkpoint_every: 10_000,
+            ..micro_cfg()
+        };
+        let summary = run_jepa_training(config);
+        let steps: Vec<usize> = summary.eval_results.iter().map(|e| e.step).collect();
+        assert_eq!(steps, vec![2, 4, 5], "cadence evals at 2,4 plus a distinct final eval at 5");
+    }
+
+    // E3. All recorded accuracies are within [0.0, 1.0].
+    #[test]
+    fn test_eval_accuracies_in_range() {
+        let config = JepaRunConfig {
+            total_steps: 4,
+            eval_every: 2,
+            log_every: 100,
+            checkpoint_every: 10_000,
+            ..micro_cfg()
+        };
+        let summary = run_jepa_training(config);
+        assert!(!summary.eval_results.is_empty(), "expected at least one eval result");
+        for e in &summary.eval_results {
+            assert!((0.0..=1.0).contains(&e.knn_acc), "knn_acc out of range: {}", e.knn_acc);
+            assert!((0.0..=1.0).contains(&e.probe_acc), "probe_acc out of range: {}", e.probe_acc);
+        }
+        assert!((0.0..=1.0).contains(&summary.final_knn_acc));
+        assert!((0.0..=1.0).contains(&summary.final_probe_acc));
+    }
+
+    // E4. Summary carries eval results through correctly: final_knn_acc /
+    // final_probe_acc match the last entry in eval_results.
+    #[test]
+    fn test_summary_final_eval_matches_last_entry() {
+        let config = JepaRunConfig {
+            total_steps: 4,
+            eval_every: 2,
+            log_every: 100,
+            checkpoint_every: 10_000,
+            ..micro_cfg()
+        };
+        let summary = run_jepa_training(config);
+        let last = summary.eval_results.last().expect("expected at least one eval result");
+        assert_eq!(summary.final_knn_acc, last.knn_acc);
+        assert_eq!(summary.final_probe_acc, last.probe_acc);
+    }
+
+    // E5. eval_every = 0 disables evaluation entirely (no cadence, no final).
+    #[test]
+    fn test_eval_disabled_when_zero() {
+        let config = JepaRunConfig {
+            total_steps: 4,
+            eval_every: 0,
+            log_every: 100,
+            checkpoint_every: 10_000,
+            ..micro_cfg()
+        };
+        let summary = run_jepa_training(config);
+        assert!(summary.eval_results.is_empty(), "eval_every=0 must disable evaluation");
+        assert_eq!(summary.final_knn_acc, 0.0);
+        assert_eq!(summary.final_probe_acc, 0.0);
+    }
 }

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

@@ -791,6 +791,19 @@ impl JepaTrainerV2 {
         self.context_encoder.as_any_mut().downcast_mut::<CpuViTEncoder>()
     }
 
+    /// Host-side view of the context encoder weights regardless of backend:
+    /// the `CpuViTEncoder` itself, or the host copy inside a `GpuViTEncoder`
+    /// (whose device buffers are uploaded from it). Used for checkpointing.
+    pub fn context_encoder_cpu_weights(&self) -> Option<&CpuViTEncoder> {
+        if let Some(cpu) = self.context_encoder_as_cpu() {
+            return Some(cpu);
+        }
+        self.context_encoder
+            .as_any()
+            .downcast_ref::<super::jepa_gpu::GpuViTEncoder>()
+            .map(|g| g.cpu_weights())
+    }
+
     /// Sample simple random context / target patch split.
     /// Returns (context_indices, target_indices).
     fn sample_mask(&mut self, num_patches: usize) -> (Vec<usize>, Vec<usize>) {

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

@@ -334,7 +334,7 @@ pub use jepa_metrics::{
 };
 pub use jepa_runner::{
     JepaRunConfig, JepaTrainingSummary, JepaCheckpoint, JepaStepResult,
-    run_jepa_training, parse_config_from_str, ViTSizeStr,
+    JepaEvalResult, run_jepa_training, parse_config_from_str, ViTSizeStr,
 };
 pub use jepa_cluster::{
     GpuSpec, NodeSpec, ClusterTopology, FabricType,