rustytorch/docs/archive/legacy/phase2.md

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 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 —