feat: implement native FSDP full-graph caputre Phase 1 (inference) + …

[Amertos]

🚀 FSDP Full-Graph Capture – Initial Implementation Complete (Phase 1 + Phase 2)

Hey team! I've been working on the FSDP full-graph capture feature and have an alpha implementation ready. Here's what's been built:

---

✅ What's Implemented

Phase 1 – Inference Full-Graph Capture

• Full computation graph tracing that captures both PyTorch operators and FSDP communication collectives (all-gather, reduce-scatter, all-reduce) as first-class graph nodes in a unified DAG.
• FSDPInferenceCapture with:
  • torch._dynamo.export-based tracing (with manual fallback)
  • Communication-computation overlap scheduling via CUDA streams
  • DOT graph export for visualization
  • FSDPHookContext that monkey-patches torch.distributed collectives for transparent capture

Phase 2 – Training Full-Graph Capture

• Extends capture to training scenarios, including backward-pass FSDP reduce-scatter collectives.
• FSDPTrainingCapture with:
  • Forward + backward graph unification into a single training graph
  • Gradient node tracking
  • AutoRecompute metadata annotation (marks nodes with >1M elements for recomputation)
  • Memory budget awareness

Graph Optimizer (6 passes)

• dead_code_elimination – removes unreachable nodes
• comm_fusion – fuses consecutive same-kind communication collectives
• compute_comm_fusion – marks adjacent compute+comm nodes for fused kernels
• comm_computation_overlap – schedules all-gather on separate CUDA streams
• auto_recompute – marks large activations for recomputation
• deterministic_align – marks all nodes for reproducible execution

Core Infrastructure

• FXGraph – unified graph IR with topological sort, DOT export
• FSDPHookContext – context manager for patching/unpatching torch.distributed collectives
• patches.py – monkey-patching for all_gather, reduce_scatter, all_reduce with recording buffer
• magicompile() – convenience API as main entry point

---

📊 Test Results

59/59 tests passing ✅ (14.09s execution time)

Test Suite	Tests	Status
test_core_graph.py	14	✅ All pass
test_fsdp_inference.py	18	✅ All pass
test_fsdp_training.py	12	✅ All pass
test_optimizations.py	11	✅ All pass
test_utils_patches.py	4	✅ All pass

---

📦 Project Structure

MAGICOMPILER/
├── pyproject.toml                    # Python package (torch>=2.12)
├── README.md                         # Full documentation
├── src/magicompiler/
│   ├── __init__.py                   # Public API exports
│   ├── core/
│   │   ├── graph.py                  # FXGraph IR (Node, Edge, NodeKind)
│   │   └── compiler.py               # MagiCompiler + magicompile()
│   ├── fsdp/
│   │   ├── hooks.py                  # FSDPHookContext (patch/unpatch)
│   │   ├── graph_capture.py          # Base tracing engine
│   │   ├── inference.py              # Phase 1: FSDP Inference
│   │   ├── training.py               # Phase 2: FSDP Training
│   │   └── optimizations.py          # DCE, fusion, overlap, recompute
│   └── utils/
│       └── patches.py                # torch.distributed monkey-patching
├── tests/                            # 5 test suites, 59 tests
└── examples/
    ├── fsdp_inference_example.py     # End-to-end inference example
    └── fsdp_training_example.py      # End-to-end training example

---

🔜 Known Gaps / Next Steps

1. torch._dynamo.export is deprecated in PyTorch 2.13 – needs migration to torch.export.export (18 FutureWarnings currently)
2. AutoRecompute is metadata-only – actual activation checkpointing logic needs to be wired up
3. Distributed tests – need torchrun-based tests with multiple GPUs to verify real FSDP collective capture
4. MegaKernel custom backend – still in research phase

---

💡 How to Test It

pip install -e .
pytest tests/ -v

# Run examples
python examples/fsdp_inference_example.py
python examples/fsdp_training_example.py --compile

I can open a PR with this if folks want to review the code. Happy to iterate on feedback! 🚀