208 lines
7.2 KiB
Markdown
208 lines
7.2 KiB
Markdown
RustyTorch++ — Phase 2 Plan (Tensor API, Autograd & Graph IR)
|
||
|
||
Objective (3–4 months): Build a usable developer surface with a Tensor API, a reverse‑mode autograd engine, and a unified Graph IR that compiles to GPU via rustg. Maintain strict determinism, expand operator coverage, and deepen profiling/CI gates.
|
||
|
||
Reference paths
|
||
|
||
RUSTYTORCH_ROOT: /home/osobh/projects/rustytorch
|
||
|
||
RUSTG_COMPILER: /home/osobh/projects/rust/rustg (RTX 5090, sm_120)
|
||
|
||
STRATOSWARM_ROOT: /home/osobh/projects/stratoswarm (deployment later)
|
||
|
||
Darwin apps are deferred. Phase 2 focuses solely on RustyTorch++ core libraries.
|
||
|
||
1. Scope & Non‑Goals
|
||
|
||
In scope
|
||
|
||
Public Tensor API (shape/dtype/device), views/strides, broadcasting, and basic indexing
|
||
|
||
Autograd engine (reverse‑mode) with tape + graph hybrid, higher‑order gradients disabled by default
|
||
|
||
Graph IR (mid‑level) with operator metadata (cost/shape), graph rewrites, and serialization
|
||
|
||
Expanded operator set with accurate forward/backward definitions (see §4)
|
||
|
||
AMP/bf16/fp16 aware autograd with loss-scaler hooks
|
||
|
||
Determinism suite for forward/backward; seed & graph capture policies
|
||
|
||
Profiler enhancements to include autograd edges and IR pass timings
|
||
|
||
Out of scope (Phase 3+)
|
||
|
||
Distributed training and FSDP/ZeRO sharding (design only in Phase 2)
|
||
|
||
Advanced autotuning/auto kernel synthesis (keep basic param sweep)
|
||
|
||
Inference scheduler & serving stack
|
||
|
||
2. Success Criteria (Exit / Phase‑Gate)
|
||
|
||
Tensor API supports creation, device transfers, views, broadcasting, and basic indexing with compile‑time safety where feasible.
|
||
|
||
Autograd computes correct gradients for the Phase 2 op set; parity vs. fp32 numeric goldens; no tape leaks.
|
||
|
||
Graph IR compiles representative model blocks (Transformer MLP/Attention forward) through rustg; IR round‑trip golden tests.
|
||
|
||
Performance: ≥ 20% reduction in end‑to‑end step time vs. Phase 1 on macrobench (forward+backward), with graph capture hit‑rate ≥ 75%.
|
||
|
||
Determinism: fixed‑seed forward/backward reproducible within tolerances (fp32 ≤ 1e‑6; bf16/fp16 ≤ 1e‑3); checkpoints resume cleanly.
|
||
|
||
Observability: profiler emits autograd edge timings; IR pass timing; memory peaks by phase (fwd/bwd/opt).
|
||
|
||
3. Architecture Work
|
||
|
||
A. Tensor API
|
||
|
||
Core Tensor<T> with shape/stride/dtype/device; type‑state for residency (HostPinned/Device)
|
||
|
||
Views: reshape/permute/expand/slice; broadcasting semantics; contiguous/non‑contiguous flags
|
||
|
||
DLPack interop (import/export); zero‑copy where safe
|
||
|
||
B. Autograd Engine
|
||
|
||
Tape of Node (op, inputs, outputs, ctx) with release strategy to avoid leaks
|
||
|
||
Backward registry: per‑op gradient fns with saved tensors / ctx
|
||
|
||
Checkpointing hooks (activation rematerialization) integrated with memory planner
|
||
|
||
AMP‑aware gradients; loss scaler callbacks for overflow recovery
|
||
|
||
C. Graph IR (mid‑level)
|
||
|
||
Nodes: ops with cost/shape metadata; edges with layout/stride annotations
|
||
|
||
Passes: canonicalize → shape infer → layout normalize → fuse (MLP/norm/attention epilogues) → schedule hints
|
||
|
||
Serialization: \*.rtxir with schema version; golden snapshots; IR → rustg lowering
|
||
|
||
4. Operator Coverage (Phase 2)
|
||
|
||
Forward + Backward implemented
|
||
|
||
Tensor creation & elementwise: add, sub, mul, div, pow, exp, log, tanh, relu, gelu
|
||
|
||
Linear algebra: matmul (batched), gemm_bias, einsum limited patterns
|
||
|
||
Reductions: sum, mean, amax, amin, var, std (stable reductions option)
|
||
|
||
Normalizations: layer_norm, rms_norm
|
||
|
||
Attention building blocks: softmax (stable), rotary embeddings (RoPE), mask ops; (Phase 3: FlashAttention full)
|
||
|
||
Tensor transforms: reshape, permute, contiguous, slice, concat, stack
|
||
|
||
Type/Device: to(dtype), to(device) (async transfers), pin_memory
|
||
|
||
5. Benchmarks & Determinism
|
||
|
||
Micro
|
||
|
||
Op‑wise forward/backward latency and throughput; gradient parity vs. analytical/finite diff (small shapes)
|
||
|
||
Macro
|
||
|
||
Transformer block (MLP + attention forward/backward) on RTX 5090 with sequence 1k–4k; report step time & memory peak
|
||
|
||
Determinism
|
||
|
||
Fixed seeds, fixed inputs; saved graphs and kernel cache keys; verify parity thresholds
|
||
|
||
6. Memory & Execution Policies
|
||
|
||
Deterministic allocation order and pooling; fragmentation histogram by phase
|
||
|
||
Overlap policy: H2D/D2H with compute for backward where possible
|
||
|
||
Checkpointing toggle with rematerialization budget; report memory vs. time trade‑off
|
||
|
||
7. Observability & Telemetry
|
||
|
||
Profiler: autograd edge timings, IR pass timings, gradient hotspot table
|
||
|
||
Metrics: step_time_ms, fwd_time_ms, bwd_time_ms, alloc_frag_pct, graph_hit_pct, dram_gbps, occ_pct
|
||
|
||
Dashboards: compare Phase 1 vs. Phase 2; per‑op gradient cost
|
||
|
||
8. CI Gates (Phase 2)
|
||
|
||
Determinism Gate: forward/backward tolerance thresholds; fixed seed; graph/hash match
|
||
|
||
Perf Gate: ≥20% step‑time reduction vs. Phase 1 baseline; no regression >5%
|
||
|
||
Memory Gate: peak ≤ budget; frag% ≤ 15%; no leaks in 8h soak
|
||
|
||
Security Gate: pinned deps; audits clean; signed artifacts & SBOM
|
||
|
||
9. Deliverables
|
||
|
||
rtx-tensor crate (public API) with views/broadcasting/indexing
|
||
|
||
rtx-autograd crate with tape engine + backward registry
|
||
|
||
rtx-ir crate (mid‑level IR) with passes and serializer; rtx‑compiler wires lowering → rustg
|
||
|
||
Expanded fused kernels for backward paths (norm/MLP epilogues)
|
||
|
||
Micro/macro benches + deterministic harness; baseline JSONs
|
||
|
||
Docs: docs/tensor_api.md, docs/autograd.md, docs/graph_ir.md; ADRs updated
|
||
|
||
10. RACI — Phase 2 Agent Ownership
|
||
|
||
Conductor (Orchestrator) — Accountable: roadmap coherence, gate readiness
|
||
|
||
Rust Engineer — Responsible: tensor storage/residency types; async transfers; API safety
|
||
|
||
Frontend Compiler Agent — Responsible: IR passes; cost/shape metadata; serializer
|
||
|
||
IR Rewrite & Fusion Agent — Responsible: fusion patterns (MLP/norm epilogues); schedule hints
|
||
|
||
Kernel Synthesizer Agent — Responsible: backward kernels; launch param tuning; fallbacks
|
||
|
||
Memory Planner Agent — Responsible: pooling policy; checkpointing integration; frag metrics
|
||
|
||
Performance Engineer — Accountable: fwd/bwd benches; profiler enhancements; perf gates
|
||
|
||
Security Engineer & Auditor — Accountable: audits; SBOM/sign
|
||
|
||
DevOps Engineer — Responsible: CI matrix; GPU runners; artifact storage
|
||
|
||
Governance & Provenance Agent — Responsible: provenance manifests; signatures
|
||
|
||
Agent Organizer — Accountable: tool allowlists; routing policies
|
||
|
||
11. Risks & Mitigations
|
||
|
||
Gradient correctness drift → finite‑diff checks on small tensors; analytical goldens; nightly parity jobs
|
||
|
||
Autograd memory bloat → aggressive release policy; checkpointing; saved‑tensor pruning
|
||
|
||
Broadcast semantics bugs → comprehensive shape/property tests; assert non‑contiguous invariants
|
||
|
||
IR stability → versioned schema; golden snapshots; migration notes
|
||
|
||
12. Timeline (suggested)
|
||
|
||
Weeks 1–2: Tensor API skeleton; device/types; views; DLPack interop
|
||
|
||
Weeks 3–4: Autograd tape engine; backward registry for core ops; AMP hooks
|
||
|
||
Weeks 5–6: Graph IR crate; passes; serializer; lowering paths to rustg
|
||
|
||
Weeks 7–8: Backward fused kernels; benchmarks; profiler updates; CI gates online; phase‑gate review
|
||
|
||
13. Phase‑3 Handover Seeds
|
||
|
||
Distributed training design draft: NCCL/RCCL primitives, topology discovery, overlap policies
|
||
|
||
Hybrid parallel prototypes: tensor/pipeline/data split APIs
|
||
|
||
Early FSDP/ZeRO notes: partition strategies; optimizer sharding API
|
||
|
||
— End of Phase 2 Plan —
|