feat(jepa): GPU weight re-upload on checkpoint resume

GpuViTEncoder gains upload_weights (extracted from construction),
cpu_weights_mut, and reupload_weights; JepaTrainerV2 exposes
context_encoder_as_any_mut for backend-specific downcasts. The runner
resume path now restores checkpoint fields into the GPU encoder's host
copy and pushes them back to the device buffers — previously GPU
resume restored only the step counter with a warning. If re-upload
fails after host restore, the run aborts rather than training on stale
device weights.

Verified live on the RTX 5060 Ti: train 20 steps -> resume from the
.jepa binary with total_steps=30 -> "Resumed from step 20", exactly 10
further steps, eval runs, no warnings. New tests: GPU output changes
after host mutation + re-upload; CPU-target re-upload is a no-op Ok.
972 CPU tests / 37 GPU jepa_gpu tests pass.

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

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

@@ -461,6 +461,78 @@ impl GpuViTEncoder {
         }
     }
 
+
+    /// Upload all encoder weights from `cpu` to the device via `stream`.
+    /// Returns `None` (with a warning) if any transfer fails.
+    #[cfg(feature = "cuda")]
+    fn upload_weights(
+        stream: &std::sync::Arc<cudarc::driver::CudaStream>,
+        cpu: &CpuViTEncoder,
+    ) -> Option<GpuWeightBuffers> {
+        let patch_embed = stream.clone_htod(&cpu.patch_embed).ok()?;
+        let proj_w = stream.clone_htod(&cpu.proj_w).ok()?;
+        let proj_b = stream.clone_htod(&cpu.proj_b).ok()?;
+
+        let mut qkv_bufs = Vec::new();
+        let mut qkv_b_bufs = Vec::new();
+        let mut out_bufs = Vec::new();
+        let mut out_b_bufs = Vec::new();
+        let mut ffn1_bufs = Vec::new();
+        let mut ffn1_b_bufs = Vec::new();
+        let mut ffn2_bufs = Vec::new();
+        let mut ffn2_b_bufs = Vec::new();
+
+        for block in &cpu.blocks {
+            qkv_bufs.push(stream.clone_htod(&block.qkv_w).ok()?);
+            qkv_b_bufs.push(stream.clone_htod(&block.qkv_b).ok()?);
+            out_bufs.push(stream.clone_htod(&block.out_w).ok()?);
+            out_b_bufs.push(stream.clone_htod(&block.out_b).ok()?);
+            ffn1_bufs.push(stream.clone_htod(&block.ffn1_w).ok()?);
+            ffn1_b_bufs.push(stream.clone_htod(&block.ffn1_b).ok()?);
+            ffn2_bufs.push(stream.clone_htod(&block.ffn2_w).ok()?);
+            ffn2_b_bufs.push(stream.clone_htod(&block.ffn2_b).ok()?);
+        }
+
+        Some(GpuWeightBuffers {
+            patch_embed,
+            proj_w,
+            proj_b,
+            block_qkv_w: qkv_bufs,
+            block_qkv_b: qkv_b_bufs,
+            block_out_w: out_bufs,
+            block_out_b: out_b_bufs,
+            block_ffn1_w: ffn1_bufs,
+            block_ffn1_b: ffn1_b_bufs,
+            block_ffn2_w: ffn2_bufs,
+            block_ffn2_b: ffn2_b_bufs,
+        })
+    }
+
+    /// Mutable host-side weights. After mutating (e.g. checkpoint restore),
+    /// call [`GpuViTEncoder::reupload_weights`] so the device buffers match.
+    pub fn cpu_weights_mut(&mut self) -> &mut CpuViTEncoder {
+        &mut self.cpu_encoder
+    }
+
+    /// Re-upload the host-side weights to the GPU, replacing the existing
+    /// device buffers. No-op success on CPU-target encoders.
+    pub fn reupload_weights(&mut self) -> Result<(), String> {
+        #[cfg(feature = "cuda")]
+        {
+            if let Some(ctx) = &self.ctx {
+                let stream = ctx.default_stream();
+                match Self::upload_weights(&stream, &self.cpu_encoder) {
+                    Some(bufs) => {
+                        self.gpu_weights = Some(bufs);
+                        return Ok(());
+                    }
+                    None => return Err("GPU weight re-upload failed".to_string()),
+                }
+            }
+        }
+        Ok(())
+    }
+
     /// Create targeting CUDA device `device_id`.
     ///
     /// Under `cuda` feature:
@@ -1780,6 +1852,38 @@ mod tests {
         check(tiny_cfg(), &[0, 5, 10, 15]);
     }
 
+    // ── Weight re-upload: mutated host weights take effect on the GPU ────────
+    #[cfg(feature = "cuda")]
+    #[test]
+    fn test_reupload_weights_changes_gpu_output() {
+        let cfg = mini_cfg();
+        let mut gpu_enc = GpuViTEncoder::cuda(cfg.clone(), 0);
+        assert!(gpu_enc.has_gpu_weights());
+
+        let patch_indices = [0usize, 3];
+        let before = gpu_enc.encode(&patch_indices);
+
+        // Mutate host weights (scale the patch projection), then re-upload.
+        for w in &mut gpu_enc.cpu_weights_mut().proj_w {
+            *w *= 2.0;
+        }
+        gpu_enc.reupload_weights().expect("re-upload must succeed");
+
+        let after = gpu_enc.encode(&patch_indices);
+        assert_eq!(before.len(), after.len());
+        assert!(
+            before.iter().zip(after.iter()).any(|(a, b)| (a - b).abs() > 1e-6),
+            "GPU output must change after host weight mutation + re-upload"
+        );
+    }
+
+    // Re-upload is a no-op success on a CPU-target encoder (no device).
+    #[test]
+    fn test_reupload_weights_cpu_target_is_ok() {
+        let mut enc = GpuViTEncoder::cpu(mini_cfg());
+        assert!(enc.reupload_weights().is_ok());
+    }
+
     // ── 35. Full-depth numerical parity (depth=12, embed_dim=192) ────────────
     //
     // `GpuViTEncoder::try_encode_gpu` now runs *every* configured block

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

@@ -929,17 +929,27 @@ 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 if let Some(gpu_enc) = trainer
+                    .context_encoder_as_any_mut()
+                    .downcast_mut::<super::jepa_gpu::GpuViTEncoder>()
+                {
+                    // Restore into the GPU encoder's host copy, then push the
+                    // updated weights back to the device buffers.
+                    let _ = apply_fields_to_encoder(gpu_enc.cpu_weights_mut(), &ckpt.fields);
+                    if let Err(e) = gpu_enc.reupload_weights() {
+                        eprintln!(
+                            "Warning: checkpoint restored on host but GPU re-upload \
+                             failed ({e}); aborting resume to avoid training on stale \
+                             device weights"
+                        );
+                        std::process::exit(1);
+                    }
                 } else {
                     eprintln!(
-                        "Warning: resume with a GPU encoder restores the step counter \
+                        "Warning: resume restores the step counter but not model \
-                         but not model weights (GPU weight re-upload not implemented)"
+                         weights (unknown encoder backend)"
                     );
                 }
                 eprintln!("Resumed from step {}", ckpt.step);

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

@@ -791,6 +791,12 @@ impl JepaTrainerV2 {
         self.context_encoder.as_any_mut().downcast_mut::<CpuViTEncoder>()
     }
 
+    /// Mutable `Any` access to the context encoder, for backend-specific
+    /// downcasts (e.g. `GpuViTEncoder` weight restore + re-upload).
+    pub fn context_encoder_as_any_mut(&mut self) -> &mut dyn std::any::Any {
+        self.context_encoder.as_any_mut()
+    }
+
     /// 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.