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]>
This commit is contained in:
Omar Sobh
2026-06-27 14:48:31 +00:00
co-authored by Claude Sonnet 4.6
parent 25322b019d
commit 1c822120d3
5 changed files with 3350 additions and 0 deletions
@@ -15,6 +15,51 @@ pub struct ViTConfig {
}
impl ViTConfig {
/// ViT-Tiny/16 configuration
pub fn tiny() -> Self {
Self {
image_size: 224,
patch_size: 16,
num_classes: 1000,
embed_dim: 192,
depth: 12,
num_heads: 3,
mlp_ratio: 4.0,
dropout: 0.0,
attention_dropout: 0.0,
}
}
/// ViT-Small/16 configuration
pub fn small() -> Self {
Self {
image_size: 224,
patch_size: 16,
num_classes: 1000,
embed_dim: 384,
depth: 12,
num_heads: 6,
mlp_ratio: 4.0,
dropout: 0.0,
attention_dropout: 0.0,
}
}
/// ViT-Small/14 configuration
pub fn small_14() -> Self {
Self {
image_size: 224,
patch_size: 14,
num_classes: 1000,
embed_dim: 384,
depth: 12,
num_heads: 6,
mlp_ratio: 4.0,
dropout: 0.0,
attention_dropout: 0.0,
}
}
/// ViT-Base/16 configuration
pub fn base_16() -> Self {
Self {
@@ -45,6 +90,21 @@ impl ViTConfig {
}
}
/// ViT-Large/14 configuration (alternative to large_16)
pub fn large_14() -> Self {
Self {
image_size: 224,
patch_size: 14,
num_classes: 1000,
embed_dim: 1024,
depth: 24,
num_heads: 16,
mlp_ratio: 4.0,
dropout: 0.0,
attention_dropout: 0.0,
}
}
/// ViT-Huge/14 configuration
pub fn huge_14() -> Self {
Self {