143 lines
5.6 KiB
Markdown
143 lines
5.6 KiB
Markdown
# RustyTorch++ Next Steps - December 28, 2025
|
|
|
|
## Completed Today
|
|
|
|
### TIER 0 (Critical Blockers) - DONE
|
|
- [x] Fix gradient computation in `rtx-autograd/src/func.rs`
|
|
- [x] Implement DTensor communication in `rtx-distributed/src/dtensor.rs`
|
|
- [x] Implement compiled backward ops in `rtx-autograd/src/compiled.rs`
|
|
|
|
### TIER 1 (Production Training) - DONE
|
|
- [x] INT4 quantized matmul kernels (`rtx-tensor/src/mobile/quantized.rs`)
|
|
- [x] cuSPARSELt FFI bindings (`rtx-tensor/src/cusparelt/mod.rs`)
|
|
- [x] 1F1B Pipeline executor with StageModule (`rtx-distributed/src/pipeline_parallel.rs`)
|
|
|
|
**Commits:**
|
|
- `c52ba928` - TIER 1 implementation (rebased to `744a63fa`)
|
|
|
|
---
|
|
|
|
## TIER 2 (Feature Completeness) - DONE
|
|
|
|
### 1. FlashAttention CPU/Metal Backward (~400 lines) - DONE
|
|
**Files:** `crates/training/rtx-flash-attention/src/core.rs`, `src/lib.rs`
|
|
- [x] Implement CPU backward pass (`naive_attention_backward`)
|
|
- [x] Implement Metal backward pass (wired trait impl to kernel)
|
|
- [x] Add gradient tests (30 tests pass)
|
|
|
|
### 2. Nested vmap Support (~300 lines) - DONE
|
|
**Files:** `crates/core/rtx-autograd/src/vmap.rs`
|
|
- [x] Track batch dimensions through nested calls (`BatchDimStack`)
|
|
- [x] Handle dimension collapsing/expansion (`MultiBatchedVariable`)
|
|
- [x] Add level tracking with RAII guards (`VmapLevelGuard`)
|
|
- [x] Tests pass (12 vmap tests)
|
|
|
|
### 3. Hessian-Vector Product (hvp) (~200 lines) - DONE
|
|
**Files:** `crates/core/rtx-autograd/src/func.rs`
|
|
- [x] Implement `hvp(f, primals, tangents)` with forward-over-reverse mode
|
|
- [x] Implement `hvp_finite_diff` for validation
|
|
- [x] Implement `vhp` (vector-Hessian product)
|
|
- [x] Fixed f32 precision issues (use f64 intermediate, larger epsilon 1e-3)
|
|
- [x] Add numerical validation tests (7 hvp tests pass)
|
|
|
|
### 4. Dynamic Shape Guards (~500 lines) - DONE
|
|
**Files:** `crates/specialized/rtx-synthesis/src/aot_impl/shape_guards.rs`, `src/aot.rs`
|
|
- [x] Implement shape guard generation (`generate_guards_from_operations`, `generate_symbolic_guards`)
|
|
- [x] Add recompilation triggers for shape changes (`GuardCheckResult`, `GuardFailure`)
|
|
- [x] Cache compiled graphs by shape signature (`ShapeGuardManager`, `ShapeSignature`)
|
|
- [x] Shape dimension types: Concrete, Symbolic, Bounded, Dynamic
|
|
- [x] Shape bucketing for cache efficiency
|
|
- [x] 20 tests passing
|
|
|
|
### 5. Advanced Quantization (AWQ, GPTQ, SmoothQuant) (~1050 lines) - DONE
|
|
**Files:** `crates/training/rtx-compress/src/quantization/advanced.rs`
|
|
- [x] AWQ (Activation-aware Weight Quantization) - per-channel scales, group quantization
|
|
- [x] GPTQ (Accurate Post-Training Quantization) - Hessian-based, block-wise quantization
|
|
- [x] SmoothQuant (migration difficulty from activations to weights) - configurable alpha
|
|
- [x] QuantizedTensorData dequantization support
|
|
- [x] 9 tests passing
|
|
|
|
---
|
|
|
|
## TIER 3 (Nice to Have) - DONE
|
|
|
|
### 1. Distributed Checkpoint (DCP) (~850 lines) - DONE
|
|
**Files:** `crates/training/rtx-distributed/src/dcp.rs`
|
|
- [x] Async checkpointing (AsyncSaveHandle with progress tracking)
|
|
- [x] Sharded state dict save/load (per-rank parallel I/O)
|
|
- [x] Resumption from partial checkpoints (configurable min_shards_for_partial)
|
|
- [x] World size change handling (shard redistribution)
|
|
- [x] Atomic writes with fsync
|
|
- [x] 11 tests passing
|
|
|
|
### 2. Context Parallel (~550 lines) - DONE
|
|
**Files:** `crates/training/rtx-distributed/src/context_parallel.rs`
|
|
- [x] Sequence dimension parallelism (SequenceShardInfo with even/uneven splits)
|
|
- [x] Ring attention integration (RingAttentionState, ring_attention method)
|
|
- [x] Long-context training support (max 128K tokens)
|
|
- [x] KV cache distribution across CP ranks
|
|
- [x] Async KV prefetch configuration
|
|
- [x] 11 tests passing
|
|
|
|
### 3. Autograd Profiler (~800 lines) - DONE
|
|
**Files:** `crates/core/rtx-autograd/src/profiler.rs`
|
|
- [x] Operation timing (ProfiledEvent with duration tracking)
|
|
- [x] Memory tracking per op (MemorySnapshot, allocation/deallocation)
|
|
- [x] Gradient flow visualization (GradientFlow, DOT export)
|
|
- [x] Bottleneck detection (BottleneckInfo, severity analysis)
|
|
- [x] Chrome trace export for visualization
|
|
- [x] RAII RecordGuard for scoped profiling
|
|
- [x] 12 tests passing
|
|
|
|
### 4. SDPA Backend Auto-Selection (~700 lines) - DONE
|
|
**Files:** `crates/training/rtx-flash-attention/src/backend_selector.rs`
|
|
- [x] Automatic FlashAttention vs Math vs Memory-efficient selection
|
|
- [x] Hardware detection for optimal backend (HardwareCapabilities)
|
|
- [x] Fallback chain management (alternatives ranking)
|
|
- [x] Performance-based scoring (sequence length, memory, speedup)
|
|
- [x] Auto-tuning with performance history
|
|
- [x] Debug mode and preferred backend options
|
|
- [x] 12 tests passing
|
|
|
|
---
|
|
|
|
## Known Issues to Address
|
|
|
|
### Autograd
|
|
- `hessian()` is a stub - depends on working `grad()`
|
|
- `functionalize` and `linearize` not implemented
|
|
- Anomaly detection not implemented
|
|
|
|
### Distributed
|
|
- Process group caching missing in DeviceMesh
|
|
- Async checkpointing not implemented
|
|
- FSDP2 communication ops need real NCCL wiring
|
|
|
|
### Quantization
|
|
- MX GPU kernels fall back to CPU (no CUDA kernels yet)
|
|
- QLinear/QConv quantized ops not implemented
|
|
- FP8 training pipeline incomplete (no loss scaling)
|
|
|
|
### Sparsity
|
|
- cuSPARSELt FFI bindings exist but no actual CUDA library calls
|
|
- Sparse backward (gradients through sparse tensors) not implemented
|
|
- Movement pruning not implemented
|
|
|
|
---
|
|
|
|
## Performance Benchmarks Needed
|
|
|
|
1. INT4 matmul vs FP16 baseline
|
|
2. 2:4 sparsity speedup on A100/H100
|
|
3. Pipeline parallelism scaling (1→8 GPUs)
|
|
4. FSDP2 memory reduction vs FSDP1
|
|
|
|
---
|
|
|
|
## Documentation TODO
|
|
|
|
- [ ] Add PyTorch migration guide
|
|
- [ ] Create quickstart for each domain (vision, NLP, etc.)
|
|
- [ ] Document INT4/FP8 quantization workflows
|
|
- [ ] Add model zoo with pre-trained weights
|