Commit Graph
6 Commits
Author SHA1 Message Date
osobhandClaude Sonnet 5 4aaa36a57a style: cargo fmt --workspace (whitespace/wrapping only, no semantic change)
Whole-workspace rustfmt pass picked up while iterating on Mamba GPU
backward work. Verified formatting-only via diff sampling; no logic
changed.

Co-Authored-By: Claude Sonnet 5 <[email protected]>
2026-08-10 07:09: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 e1b4061c23 feat(jepa): extended GPU training, data pipeline, integration, and cargo config
CI / Format Check (push) Failing after 12s
CI / Build (macos-latest) (push) Failing after 12s
CI / Build (ubuntu-latest) (push) Failing after 19s
CI / Distributed Training Tests (push) Has been skipped
CI / Clippy Check (push) Failing after 19s
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 8s
CI / Build CPU-Only (Explicit) (push) Failing after 16s
Documentation / Build API Documentation (push) Failing after 13s
CI / CI Success (push) Failing after 0s
Performance Benchmarks / Run Benchmarks (push) Successful in 43s
Extends jepa_train with distributed launcher, jepa_data with advanced
sampling and preprocessing, jepa_gpu with full CUDA kernel wiring,
jepa_distributed/runner/metrics/vit with additional training stages.
Adds jepa_integration module and project-local cargo config.

Co-Authored-By: Claude Sonnet 4.6 <[email protected]>
2026-06-29 21:35:34 +00: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
Omar SobhandClaude Sonnet 4.6 1c822120d3 feat(batch24-26): JEPA ViT wiring, data pipeline, cluster-scale config
CI / Build (macos-latest) (push) Failing after 30s
CI / Distributed Training Tests (push) Has been skipped
CI / Build CPU-Only (Explicit) (push) Failing after 2m20s
Documentation / Build User Guide (push) Successful in 7s
CI / Format Check (push) Failing after 16s
CI / Clippy Check (push) Failing after 1m6s
CI / Build (ubuntu-latest) (push) Failing after 1m3s
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 API Documentation (push) Failing after 2m3s
CI / CI Success (push) Failing after 0s
Performance Benchmarks / Run Benchmarks (push) Successful in 7m58s
Batch 24 — Real ViT encoder integration (ssl/jepa_vit.rs):
- JepaEncoder trait: encode(patch_indices) + embed_dim + num_patches
- CpuViTEncoder: sinusoidal+learned pos embed, LCG-init weights, GELU FFN,
  MHSA scaled dot-product; runs min(depth,2) blocks for CPU test speed
- EmaViTEncoder: shadow weights, tau-weighted update, τ=1 frozen / τ=0 copy
- JepaTrainerV2: mask→CpuViTEncoder→predictor→EmaViT→L2→EMA; timing metrics
- JepaViTConfig: tiny/small/base/large/huge presets (embed_dim, depth, heads)
- 46 tests

Batch 24b — ViT-S/T/small-14/large-14 configs (rtx-vision/configs.rs):
- Added ViTConfig::tiny() d=192/depth=12/heads=3
- Added ViTConfig::small() d=384/depth=12/heads=6
- Added ViTConfig::small_14() d=384/patch=14
- Added ViTConfig::large_14() d=1024/depth=24/patch=14

Batch 25 — ImageNet-scale streaming data pipeline (ssl/jepa_data.rs):
- ImageRecord: HWC pixel buffer with label and key
- MultiScaleRandomCrop: LCG PRNG + bilinear resampling, scale 0.2-1.0
- RandomHorizontalFlip: stochastic row mirror
- JepaAugmentationPipeline: crop→flip→ImageNet normalize (mean/std)
- InMemoryShard: synthetic LCG data for testing
- JepaBatch: augmented images + context/target indices per sample
- JepaDataPipeline: streaming iterator, Fisher-Yates epoch shuffle,
  next_batch() → None at epoch end, reset_epoch()
- DatasetStats: mask efficiency, avg context/target patch counts
- WebDatasetShard: filesystem shard descriptor stub (to_in_memory for tests)
- 35 tests

Batch 26 — Cluster-scale training configuration (ssl/jepa_cluster.rs):
- GpuSpec: RTX 5060 Ti (SM_120), RTX 4090, A100-80GB specs
- NodeSpec + ClusterTopology: homogeneous/heterogeneous cluster descriptors
- JepaParallelConfig: TP/PP/DP with for_model_and_cluster() auto-select
  (TP≥4 for ViT-L 300M+, TP=8/PP=2 for ViT-H 600M+)
- GradientCompressionConfig: TopK/PowerSGD/1-bit SGD with error feedback
- DcpCheckpointConfig: async save, EMA weights, keep-last-N
- JepaClusterConfig: validate(), memory_per_gpu_gb(), throughput estimate
- ClusterTrainingPlan: steps_per_epoch, total_steps, estimated_hours, summary
- AdaptiveBatchSizer: GNS-based batch doubling/halving with [min,max] clamp
- 42 tests

Total new: 163 JEPA tests (0 failures), 3,350 lines

Co-Authored-By: Claude Sonnet 4.6 <[email protected]>
2026-06-27 14:48:31 +00:00