Memory-budget-aware activation checkpoint planner for distributed training experiments.
-
Updated
Jul 28, 2026 - Python
#
activation-checkpointing
Here are 6 public repositories matching this topic...
Diagnosing non-monotonic activation memory in torch.compile's memory-budget partitioner, with a use-site rematerialization fix.
pytorch memory-optimization gradient-checkpointing ml-systems rematerialization torch-compile activation-checkpointing deep-learning-compilers
-
Updated
Oct 5, 2026 - Python
Stanford CS336 Assignment 2 (Systems), implemented from scratch: FlashAttention-2 in PyTorch + Triton (fwd + two-pass bwd), DDP variants, ZeRO-1 optimizer sharding, FSDP with mixed precision, profiling/benchmark harness, full writeup
cuda pytorch transformer triton cs336 performance-benchmarking gpu-kernels distributed-training mixed-precision distributed-data-parallel fsdp llm-training flash-attention activation-checkpointing zero-optimizer
-
Updated
Sep 14, 2026 - Python
Memory-efficient distributed LLM training from scratch: Triton FlashAttention, fused LM-head/cross-entropy, custom FSDP, and activation checkpointing. Part 2 of 2, extending Stanford CS336.
-
Updated
Oct 12, 2026 - Python
Automatic activation checkpointing and tensor swapping for DNN training using torch.fx graph profiling.
-
Updated
Aug 30, 2026 - Jupyter Notebook
Whole-graph activation checkpointing planner for torch.compile: a reimplementation of Jane Street's 2026 intern planner, with a memory model that follows Inductor, a search that gets past dead ends, and measured costs
-
Updated
Oct 12, 2026 - Python
Add this topic to your repo
To associate your repository with the activation-checkpointing topic, visit your repo's landing page and select "manage topics."