How rocm/jax-titan works
JAX-Titan trains models through a single-controller Shardy MPMD pipeline, with an SPMD baseline, pluggable operators, data loading, and checkpointing. Run records and observability support performance analysis; a separate planning branch estimates model and system budgets. The README documents several capabilities whose detailed implementation was not sampled, so those branches are represented at subsystem level. Edges rely on README statements or the supplied source index and excerpts; unsampled internals are not inferred.
Main entry points: train.py Training Entry · adapters.py Training Layouts · budget.py Budget Planner · summarize.py Run Summaries
Components 32
Training Runtime
- Training Entry
src/jax_titan/trainer/train.py - Training Layouts
src/jax_titan/trainer/adapters.py - Layout Sharding
src/jax_titan/trainer/layout.py - MPMD Pipeline
src/jax_titan/mpmd/pipeline.py - Parallel Shardings
src/jax_titan/parallel/sharding.py - Stage Checkpoints
src/jax_titan/trainer/checkpoint.py - Step Watchdog
src/jax_titan/observability/watchdog.py
Models and Data
- Llama Model
src/jax_titan/models/llama.py - MaxText Adapter
src/jax_titan/models/adapters/maxtext.py - MaxDiffusion Adapter
src/jax_titan/models/adapters/maxdiffusion.py - Text Sampling
src/jax_titan/data/text.py - Llama Tokenizer
src/jax_titan/data/tokenizer.py
Operator System
- Semantic Operators
src/jax_titan/ops/__init__.py - Provider Registry
src/jax_titan/ops/registry.py - Native JAX Ops
src/jax_titan/ops/native.py - Accelerator Providers
src/jax_titan/ops/providers - Op Conformance
src/jax_titan/ops/conformance.py
Run Observability
- Run Records
src/jax_titan/observability/records.py - Environment Capture
src/jax_titan/observability/env.py - Power Metrics
src/jax_titan/observability/power.py - FLOPs Model
src/jax_titan/observability/flops.py - Run Summaries
src/jax_titan/observability/summarize.py - Restart Analysis
src/jax_titan/observability/ettr.py - Profiles and Figures
src/jax_titan/observability/profile.py
Performance Planning
- Budget Planner
src/jax_titan/planning/budget.py - Model Estimates
src/jax_titan/planning/models.py - GPU Catalog
src/jax_titan/planning/gpus.py - Memory Planning
src/jax_titan/planning/memory.py - Collective Model
src/jax_titan/planning/collectives.py - Step-Time Model
src/jax_titan/planning/step.py - Performance Database
src/jax_titan/observability/perfdb.py
Other components
- Training User
Connections 33
- Training User → Training Entry: starts training
src/jax_titan/trainer/train.py - Training Entry → Training Layouts: selects layout
src/jax_titan/trainer/train.py - Training Layouts → MPMD Pipeline: builds train step
src/jax_titan/trainer/adapters.py - Training Layouts → Layout Sharding: uses sharding
src/jax_titan/trainer/adapters.py - MPMD Pipeline → Parallel Shardings: uses shardings
src/jax_titan/mpmd/pipeline.py - Training Entry → Stage Checkpoints: saves state
src/jax_titan/trainer/train.py - Training Entry → Step Watchdog: monitors steps
src/jax_titan/trainer/train.py - Training Layouts → Llama Model: builds stages
src/jax_titan/trainer/adapters.py - Training Layouts → MaxText Adapter: adapts model
src/jax_titan/trainer/adapters.py - Training Layouts → MaxDiffusion Adapter: adapts model
src/jax_titan/trainer/adapters.py - Training Entry → Text Sampling: loads batches
src/jax_titan/trainer/train.py - Text Sampling → Llama Tokenizer: tokenizes text
src/jax_titan/data/text.py - Llama Model → Semantic Operators: uses operators
src/jax_titan/models/llama.py - Semantic Operators → Provider Registry: dispatches ops
src/jax_titan/ops/__init__.py - Provider Registry → Native JAX Ops: falls back
src/jax_titan/ops/__init__.py - Provider Registry → Accelerator Providers: selects provider
src/jax_titan/ops/providers/triton.py - Op Conformance → Semantic Operators: checks implementations
src/jax_titan/ops/conformance.py - Training Entry → Run Records: writes events
src/jax_titan/trainer/train.py - Training Entry → Environment Capture: captures environment
src/jax_titan/trainer/train.py - Training Entry → Power Metrics: collects power
src/jax_titan/trainer/train.py - Training Entry → FLOPs Model: reports FLOPs
src/jax_titan/trainer/train.py - Run Summaries → Run Records: reads records
src/jax_titan/observability/summarize.py - Run Summaries → Power Metrics: reads energy
src/jax_titan/observability/summarize.py - Run Summaries → Restart Analysis: computes ETTR
src/jax_titan/observability/summarize.py - Training Entry → Profiles and Figures: captures traces
src/jax_titan/trainer/train.py - Training User → Budget Planner: requests estimate
README.md - Budget Planner → Model Estimates: estimates model
src/jax_titan/planning/budget.py - Budget Planner → GPU Catalog: selects GPU
src/jax_titan/planning/budget.py - Budget Planner → Memory Planning: estimates memory
src/jax_titan/planning/budget.py - Budget Planner → Collective Model: estimates communication
src/jax_titan/planning/budget.py - Budget Planner → Step-Time Model: predicts step time
README.md - Step-Time Model → Collective Model: models collectives
src/jax_titan/planning/step.py - Step-Time Model → Performance Database: checks regressions
src/jax_titan/planning/step.py