Commit Graph
6 Commits
Author SHA1 Message Date
osobhandClaude Fable 5 ac3f2af06b feat(jepa): eval in the training loop, NCCL GPU AllReduce, GPU checkpointing
CI / Format Check (push) Failing after 6s
CI / Build (ubuntu-latest) (push) Failing after 5s
Documentation / Build API Documentation (push) Failing after 7s
Documentation / Build User Guide (push) Successful in 8s
CI / Build (macos-latest) (push) Failing after 11s
CI / Test (macos-latest) (push) Has been skipped
CI / Test (ubuntu-latest) (push) Has been skipped
CI / Python Bindings (maturin) (macos-latest) (push) Has been skipped
CI / Python Bindings (maturin) (ubuntu-latest) (push) Has been skipped
CI / WASM Build + Size Check (push) Has been skipped
CI / Distributed Training Tests (push) Has been skipped
Performance Benchmarks / Run Benchmarks (push) Successful in 29s
CI / Clippy Check (push) Failing after 15s
CI / Build CPU-Only (Explicit) (push) Failing after 49s
CI / CI Success (push) Failing after 0s
- 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]>
2026-07-10 03:51:36 -07:00
osobhandClaude Fable 5 b6440905e7 feat(jepa): gzip-compressed WebDataset shard support
CI / Format Check (push) Failing after 6s
Performance Benchmarks / Run Benchmarks (push) Failing after 7s
CI / Build (ubuntu-latest) (push) Failing after 6s
CI / Clippy Check (push) Failing after 8s
Documentation / Build User Guide (push) Successful in 7s
CI / Build (macos-latest) (push) Failing after 9s
CI / Build CPU-Only (Explicit) (push) Failing after 59s
CI / CI Success (push) Failing after 0s
CI / Test (macos-latest) (push) Has been skipped
CI / Test (ubuntu-latest) (push) Has been skipped
CI / Python Bindings (maturin) (macos-latest) (push) Has been skipped
CI / Python Bindings (maturin) (ubuntu-latest) (push) Has been skipped
CI / WASM Build + Size Check (push) Has been skipped
CI / Distributed Training Tests (push) Has been skipped
Documentation / Build API Documentation (push) Failing after 16s
read_webdataset_shard detects the gzip magic bytes (1F 8B, not
extension) and decompresses via flate2 before tar parsing;
WebDatasetShard::load no longer rejects .tar.gz/.tgz. Round-trip test
writes a real gzipped tar and loads it back.

Co-Authored-By: Claude Fable 5 <[email protected]>
2026-07-09 22:37:05 -07:00
osobhandClaude Fable 5 1e3c604896 feat(meta,jepa): expose GPU features through meta-crates; wire JEPA cluster plan and real shard loading
Meta-crates (Phase 2):
- rtx-core / rtx-training / rtx-inference-stack gain cuda and metal
  features threading into their sub-crates; GPU was previously
  unreachable through the user-facing bundles.
- rtx-training restores rtx-distributed (the hpc-channels blocker is
  gone) so the advertised DistributedTransformerTrainer resolves; drops
  the unused rtx-runtime dep.
- rtx-transformers drops unused rtx-backend/rtx-backend-cpu deps
  (stale comment referenced a teacher that never used them).

Never-compiled CUDA paths fixed (surfaced by the new feature wiring,
verified on RTX 5060 Ti / CUDA 13.1):
- rtx-compress build.rs: missing Path/Command/fs imports.
- rtx-flash-attention flash_decode_forward: reborrow &mut kernel args.
- rtx-transformers: rope kernel include path, cudarc 0.18 Arc<CudaModule>,
  PushKernelArg imports in jepa_gpu, edition-2024 ref patterns.
- rtx-memory: full cudarc 0.18 port (CudaContext, stream-based alloc,
  DevicePtr accessors, error enum formatting) across gpu_pinning,
  gpu_transfer, gpu_real, gpu_allocator/arena, gpu_tests.

JEPA (Phase 3):
- JepaRunConfig::apply_cluster_plan consumes ClusterTrainingPlan
  (batch size, TP/DP, world size, total steps) so jepa_cluster is no
  longer standalone dead config; ViTSizeStr::approx_params_m feeds
  JepaParallelConfig::for_model_and_cluster.
- WebDatasetShard::load reads real .tar shards from disk via the
  existing parser (gzip rejected explicitly); to_in_memory documented
  as synthetic/test-only.
- New image-decode feature actually defines the dep for the previously
  unreachable cfg(feature = "image-decode") JPEG/PNG decode path.

Co-Authored-By: Claude Fable 5 <[email protected]>
2026-07-09 19:25:51 -07:00
Omar SobhandClaude Sonnet 4.6 f487196367 feat(batch27): JEPA ViT bridge, WebDataset shard reading, training loop
CI / Format Check (push) Failing after 11s
CI / Test (macos-latest) (push) Has been skipped
CI / Test (ubuntu-latest) (push) Has been skipped
CI / Python Bindings (maturin) (macos-latest) (push) Has been skipped
CI / Python Bindings (maturin) (ubuntu-latest) (push) Has been skipped
CI / WASM Build + Size Check (push) Has been skipped
Documentation / Build User Guide (push) Successful in 9s
CI / Build CPU-Only (Explicit) (push) Failing after 1m10s
CI / CI Success (push) Failing after 0s
Documentation / Build API Documentation (push) Failing after 40s
Performance Benchmarks / Run Benchmarks (push) Successful in 7m53s
CI / Build (macos-latest) (push) Failing after 30s
CI / Build (ubuntu-latest) (push) Failing after 48s
CI / Distributed Training Tests (push) Has been skipped
CI / Clippy Check (push) Failing after 52s
Gap 2 — rtx-vision ViT bridge (jepa_vision_bridge.rs, 8 tests):
- ViT::forward_features(): patch reps without classification head
- ViT::encode_patch_indices(): shape-correct placeholder for GPU dispatch
- RtxVisionJepaEncoder implementing JepaEncoder (vision-bridge feature)
- From<&ViTConfig> for JepaViTConfig config conversion
- rtx-vision added as optional dep; vision-bridge feature gate

Gap 3 — WebDataset tar-shard reading (jepa_data.rs, +12 tests, 47 total):
- parse_tar_bytes(): pure stdlib tar parser (512-byte block format)
- read_webdataset_shard(): file reader with ShardLoadStats timing
- WebDatasetRecord: key, image_bytes, label, extension
- ShuffleBuffer: fixed-capacity reservoir sampling via LCG PRNG
- JepaDataPipeline::from_filesystem(): validates paths, loads shards, builds pipeline

Gap 5 — Training loop runner (jepa_runner.rs + examples/jepa_train.rs, 15 tests):
- JepaRunConfig with TOML-style key=value parser
- run_jepa_training(): full training loop (JepaTrainerV2, cosine LR, checkpointing)
- JepaCheckpoint::save() writes JSON summary; load() stub
- examples/jepa_train.rs: --config/--size/--steps/--dry-run CLI flags

Co-Authored-By: Claude Sonnet 4.6 <[email protected]>
2026-06-27 15:23:55 +00:00
osobhandClaude Opus 4.8 1b1ce0604a SMT E1: SetEncoderTeacher — the predictive-state oracle (time-parallel, tape-trained)
Promotes the proven attention set-encoder smoke test (PR #4,
tape_train_smoke.rs) into a reusable rtx-transformers layer for Supervised
Memory Training. The teacher maps a window of past tokens to a fixed-size
memory M via embed → self-attention → residual → mean-pool → memory
projection, with a decoder head supervised by predict-the-future MSE so that
M becomes a sufficient statistic of the past.

- Trains end-to-end on Autodiff<CpuBackend> (the gradient-correct real backend
  from PR #3), with a self-contained deterministic host-side Adam.
- Time-parallel by construction (one window → one memory, no recurrence to
  unroll) — this is the oracle whose trajectory the recurrent Mamba updater is
  later behaviorally cloned against, so the recurrent net never needs BPTT.
- Exposes named_params/set_named_params so the caller (omni-think's
  PredictiveStateTeacher facade) owns safetensors persistence + BLAKE3 sealing.
- Adds rtx-backend + rtx-backend-cpu deps (the tape needs a concrete backend).

Tests: teacher trains (loss >5x drop), encode is deterministic + fixed-size,
params round-trip. fmt + clippy(-D warnings) clean on the new module.

Co-Authored-By: Claude Opus 4.8 (1M context) <[email protected]>
2026-06-16 09:37:55 -07:00
redclawsystems 4d88dc0584 Initial commit 2026-03-04 00:08:42 +00:00