# rocm/jax-titan architecture

> Architecture diagram of the GitHub repository rocm/jax-titan, generated by GitDiagram from its file tree, README and a sample of source files (last updated 2026-11-10). It is an AI-made overview: check details against the source.

- Interactive diagram: https://gitdiagram.com/rocm/jax-titan
- Repository: https://github.com/rocm/jax-titan

## Overview

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: [`src/jax_titan/trainer/train.py`](https://github.com/rocm/jax-titan/blob/main/src/jax_titan/trainer/train.py) (Training Entry), [`src/jax_titan/trainer/adapters.py`](https://github.com/rocm/jax-titan/blob/main/src/jax_titan/trainer/adapters.py) (Training Layouts), [`src/jax_titan/planning/budget.py`](https://github.com/rocm/jax-titan/blob/main/src/jax_titan/planning/budget.py) (Budget Planner), [`src/jax_titan/observability/summarize.py`](https://github.com/rocm/jax-titan/blob/main/src/jax_titan/observability/summarize.py) (Run Summaries)

## Diagram (Mermaid)

```mermaid
flowchart TD

subgraph group_training["Training Runtime"]
  node_trainer["Training Entry<br/>[train.py]"]
  node_layouts["Training Layouts<br/>[adapters.py]"]
  node_layout_specs["Layout Sharding<br/>[layout.py]"]
  node_pipeline["MPMD Pipeline<br/>[pipeline.py]"]
  node_parallel["Parallel Shardings<br/>[sharding.py]"]
  node_checkpoint["Stage Checkpoints<br/>[checkpoint.py]"]
  node_watchdog["Step Watchdog<br/>[watchdog.py]"]
end

subgraph group_models["Models and Data"]
  node_llama["Llama Model<br/>[llama.py]"]
  node_maxtext["MaxText Adapter<br/>[maxtext.py]"]
  node_maxdiffusion["MaxDiffusion Adapter<br/>[maxdiffusion.py]"]
  node_textdata["Text Sampling<br/>[text.py]"]
  node_tokenizer["Llama Tokenizer<br/>[tokenizer.py]"]
end

subgraph group_operators["Operator System"]
  node_ops_api["Semantic Operators<br/>[__init__.py]"]
  node_op_registry["Provider Registry<br/>[registry.py]"]
  node_native_ops["Native JAX Ops<br/>[native.py]"]
  node_provider_kernels["Accelerator Providers"]
  node_conformance["Op Conformance<br/>[conformance.py]"]
end

subgraph group_observability["Run Observability"]
  node_records["Run Records<br/>[records.py]"]
  node_environment["Environment Capture<br/>[env.py]"]
  node_power["Power Metrics<br/>[power.py]"]
  node_flops["FLOPs Model<br/>[flops.py]"]
  node_summarize["Run Summaries<br/>[summarize.py]"]
  node_ettr["Restart Analysis<br/>[ettr.py]"]
  node_profiling["Profiles and Figures<br/>[profile.py]"]
end

subgraph group_planning["Performance Planning"]
  node_budget["Budget Planner<br/>[budget.py]"]
  node_model_estimates["Model Estimates<br/>[models.py]"]
  node_gpu_catalog["GPU Catalog<br/>[gpus.py]"]
  node_memory_model["Memory Planning<br/>[memory.py]"]
  node_collective_model["Collective Model<br/>[collectives.py]"]
  node_step_model["Step-Time Model<br/>[step.py]"]
  node_perfdb[("Performance Database<br/>[perfdb.py]")]
end

node_user(("Training User"))

node_user -->|"starts training"| node_trainer
node_trainer -->|"selects layout"| node_layouts
node_layouts -->|"builds train step"| node_pipeline
node_layouts -->|"uses sharding"| node_layout_specs
node_pipeline -->|"uses shardings"| node_parallel
node_trainer -->|"saves state"| node_checkpoint
node_trainer -->|"monitors steps"| node_watchdog
node_layouts -->|"builds stages"| node_llama
node_layouts -->|"adapts model"| node_maxtext
node_layouts -->|"adapts model"| node_maxdiffusion
node_trainer -->|"loads batches"| node_textdata
node_textdata -->|"tokenizes text"| node_tokenizer
node_llama -->|"uses operators"| node_ops_api
node_ops_api -->|"dispatches ops"| node_op_registry
node_op_registry -->|"falls back"| node_native_ops
node_op_registry -->|"selects provider"| node_provider_kernels
node_conformance -->|"checks implementations"| node_ops_api
node_trainer -->|"writes events"| node_records
node_trainer -->|"captures environment"| node_environment
node_trainer -->|"collects power"| node_power
node_trainer -->|"reports FLOPs"| node_flops
node_summarize -->|"reads records"| node_records
node_summarize -->|"reads energy"| node_power
node_summarize -->|"computes ETTR"| node_ettr
node_trainer -->|"captures traces"| node_profiling
node_user -->|"requests estimate"| node_budget
node_budget -->|"estimates model"| node_model_estimates
node_budget -->|"selects GPU"| node_gpu_catalog
node_budget -->|"estimates memory"| node_memory_model
node_budget -->|"estimates communication"| node_collective_model
node_budget -->|"predicts step time"| node_step_model
node_step_model -->|"models collectives"| node_collective_model
node_step_model -->|"checks regressions"| node_perfdb

click node_trainer "https://github.com/rocm/jax-titan/blob/main/src/jax_titan/trainer/train.py"
click node_layouts "https://github.com/rocm/jax-titan/blob/main/src/jax_titan/trainer/adapters.py"
click node_layout_specs "https://github.com/rocm/jax-titan/blob/main/src/jax_titan/trainer/layout.py"
click node_pipeline "https://github.com/rocm/jax-titan/blob/main/src/jax_titan/mpmd/pipeline.py"
click node_parallel "https://github.com/rocm/jax-titan/blob/main/src/jax_titan/parallel/sharding.py"
click node_checkpoint "https://github.com/rocm/jax-titan/blob/main/src/jax_titan/trainer/checkpoint.py"
click node_watchdog "https://github.com/rocm/jax-titan/blob/main/src/jax_titan/observability/watchdog.py"
click node_llama "https://github.com/rocm/jax-titan/blob/main/src/jax_titan/models/llama.py"
click node_maxtext "https://github.com/rocm/jax-titan/blob/main/src/jax_titan/models/adapters/maxtext.py"
click node_maxdiffusion "https://github.com/rocm/jax-titan/blob/main/src/jax_titan/models/adapters/maxdiffusion.py"
click node_textdata "https://github.com/rocm/jax-titan/blob/main/src/jax_titan/data/text.py"
click node_tokenizer "https://github.com/rocm/jax-titan/blob/main/src/jax_titan/data/tokenizer.py"
click node_ops_api "https://github.com/rocm/jax-titan/blob/main/src/jax_titan/ops/__init__.py"
click node_op_registry "https://github.com/rocm/jax-titan/blob/main/src/jax_titan/ops/registry.py"
click node_native_ops "https://github.com/rocm/jax-titan/blob/main/src/jax_titan/ops/native.py"
click node_provider_kernels "https://github.com/rocm/jax-titan/tree/main/src/jax_titan/ops/providers"
click node_conformance "https://github.com/rocm/jax-titan/blob/main/src/jax_titan/ops/conformance.py"
click node_records "https://github.com/rocm/jax-titan/blob/main/src/jax_titan/observability/records.py"
click node_environment "https://github.com/rocm/jax-titan/blob/main/src/jax_titan/observability/env.py"
click node_power "https://github.com/rocm/jax-titan/blob/main/src/jax_titan/observability/power.py"
click node_flops "https://github.com/rocm/jax-titan/blob/main/src/jax_titan/observability/flops.py"
click node_summarize "https://github.com/rocm/jax-titan/blob/main/src/jax_titan/observability/summarize.py"
click node_ettr "https://github.com/rocm/jax-titan/blob/main/src/jax_titan/observability/ettr.py"
click node_profiling "https://github.com/rocm/jax-titan/blob/main/src/jax_titan/observability/profile.py"
click node_budget "https://github.com/rocm/jax-titan/blob/main/src/jax_titan/planning/budget.py"
click node_model_estimates "https://github.com/rocm/jax-titan/blob/main/src/jax_titan/planning/models.py"
click node_gpu_catalog "https://github.com/rocm/jax-titan/blob/main/src/jax_titan/planning/gpus.py"
click node_memory_model "https://github.com/rocm/jax-titan/blob/main/src/jax_titan/planning/memory.py"
click node_collective_model "https://github.com/rocm/jax-titan/blob/main/src/jax_titan/planning/collectives.py"
click node_step_model "https://github.com/rocm/jax-titan/blob/main/src/jax_titan/planning/step.py"
click node_perfdb "https://github.com/rocm/jax-titan/blob/main/src/jax_titan/observability/perfdb.py"

classDef toneNeutral fill:#f8fafc,stroke:#334155,stroke-width:1.5px,color:#0f172a
classDef toneBlue fill:#dbeafe,stroke:#2563eb,stroke-width:1.5px,color:#172554
classDef toneAmber fill:#fef3c7,stroke:#d97706,stroke-width:1.5px,color:#78350f
classDef toneMint fill:#dcfce7,stroke:#16a34a,stroke-width:1.5px,color:#14532d
classDef toneRose fill:#ffe4e6,stroke:#e11d48,stroke-width:1.5px,color:#881337
classDef toneIndigo fill:#e0e7ff,stroke:#4f46e5,stroke-width:1.5px,color:#312e81
classDef toneTeal fill:#ccfbf1,stroke:#0f766e,stroke-width:1.5px,color:#134e4a
class node_trainer,node_layouts,node_layout_specs,node_pipeline,node_parallel,node_checkpoint,node_watchdog,node_user toneBlue
class node_llama,node_maxtext,node_maxdiffusion,node_textdata,node_tokenizer toneAmber
class node_ops_api,node_op_registry,node_native_ops,node_provider_kernels,node_conformance toneMint
class node_records,node_environment,node_power,node_flops,node_summarize,node_ettr,node_profiling toneRose
class node_budget,node_model_estimates,node_gpu_catalog,node_memory_model,node_collective_model,node_step_model,node_perfdb toneIndigo
```

## Components

### Training Runtime

- [Training Entry](https://github.com/rocm/jax-titan/blob/main/src/jax_titan/trainer/train.py): `src/jax_titan/trainer/train.py`
- [Training Layouts](https://github.com/rocm/jax-titan/blob/main/src/jax_titan/trainer/adapters.py): `src/jax_titan/trainer/adapters.py`
- [Layout Sharding](https://github.com/rocm/jax-titan/blob/main/src/jax_titan/trainer/layout.py): `src/jax_titan/trainer/layout.py`
- [MPMD Pipeline](https://github.com/rocm/jax-titan/blob/main/src/jax_titan/mpmd/pipeline.py): `src/jax_titan/mpmd/pipeline.py`
- [Parallel Shardings](https://github.com/rocm/jax-titan/blob/main/src/jax_titan/parallel/sharding.py): `src/jax_titan/parallel/sharding.py`
- [Stage Checkpoints](https://github.com/rocm/jax-titan/blob/main/src/jax_titan/trainer/checkpoint.py): `src/jax_titan/trainer/checkpoint.py`
- [Step Watchdog](https://github.com/rocm/jax-titan/blob/main/src/jax_titan/observability/watchdog.py): `src/jax_titan/observability/watchdog.py`

### Models and Data

- [Llama Model](https://github.com/rocm/jax-titan/blob/main/src/jax_titan/models/llama.py): `src/jax_titan/models/llama.py`
- [MaxText Adapter](https://github.com/rocm/jax-titan/blob/main/src/jax_titan/models/adapters/maxtext.py): `src/jax_titan/models/adapters/maxtext.py`
- [MaxDiffusion Adapter](https://github.com/rocm/jax-titan/blob/main/src/jax_titan/models/adapters/maxdiffusion.py): `src/jax_titan/models/adapters/maxdiffusion.py`
- [Text Sampling](https://github.com/rocm/jax-titan/blob/main/src/jax_titan/data/text.py): `src/jax_titan/data/text.py`
- [Llama Tokenizer](https://github.com/rocm/jax-titan/blob/main/src/jax_titan/data/tokenizer.py): `src/jax_titan/data/tokenizer.py`

### Operator System

- [Semantic Operators](https://github.com/rocm/jax-titan/blob/main/src/jax_titan/ops/__init__.py): `src/jax_titan/ops/__init__.py`
- [Provider Registry](https://github.com/rocm/jax-titan/blob/main/src/jax_titan/ops/registry.py): `src/jax_titan/ops/registry.py`
- [Native JAX Ops](https://github.com/rocm/jax-titan/blob/main/src/jax_titan/ops/native.py): `src/jax_titan/ops/native.py`
- [Accelerator Providers](https://github.com/rocm/jax-titan/tree/main/src/jax_titan/ops/providers): `src/jax_titan/ops/providers`
- [Op Conformance](https://github.com/rocm/jax-titan/blob/main/src/jax_titan/ops/conformance.py): `src/jax_titan/ops/conformance.py`

### Run Observability

- [Run Records](https://github.com/rocm/jax-titan/blob/main/src/jax_titan/observability/records.py): `src/jax_titan/observability/records.py`
- [Environment Capture](https://github.com/rocm/jax-titan/blob/main/src/jax_titan/observability/env.py): `src/jax_titan/observability/env.py`
- [Power Metrics](https://github.com/rocm/jax-titan/blob/main/src/jax_titan/observability/power.py): `src/jax_titan/observability/power.py`
- [FLOPs Model](https://github.com/rocm/jax-titan/blob/main/src/jax_titan/observability/flops.py): `src/jax_titan/observability/flops.py`
- [Run Summaries](https://github.com/rocm/jax-titan/blob/main/src/jax_titan/observability/summarize.py): `src/jax_titan/observability/summarize.py`
- [Restart Analysis](https://github.com/rocm/jax-titan/blob/main/src/jax_titan/observability/ettr.py): `src/jax_titan/observability/ettr.py`
- [Profiles and Figures](https://github.com/rocm/jax-titan/blob/main/src/jax_titan/observability/profile.py): `src/jax_titan/observability/profile.py`

### Performance Planning

- [Budget Planner](https://github.com/rocm/jax-titan/blob/main/src/jax_titan/planning/budget.py): `src/jax_titan/planning/budget.py`
- [Model Estimates](https://github.com/rocm/jax-titan/blob/main/src/jax_titan/planning/models.py): `src/jax_titan/planning/models.py`
- [GPU Catalog](https://github.com/rocm/jax-titan/blob/main/src/jax_titan/planning/gpus.py): `src/jax_titan/planning/gpus.py`
- [Memory Planning](https://github.com/rocm/jax-titan/blob/main/src/jax_titan/planning/memory.py): `src/jax_titan/planning/memory.py`
- [Collective Model](https://github.com/rocm/jax-titan/blob/main/src/jax_titan/planning/collectives.py): `src/jax_titan/planning/collectives.py`
- [Step-Time Model](https://github.com/rocm/jax-titan/blob/main/src/jax_titan/planning/step.py): `src/jax_titan/planning/step.py`
- [Performance Database](https://github.com/rocm/jax-titan/blob/main/src/jax_titan/observability/perfdb.py): `src/jax_titan/observability/perfdb.py`

### Other components

- **Training User**

## Connections

- Training User → Training Entry: starts training (evidence: `src/jax_titan/trainer/train.py`)
- Training Entry → Training Layouts: selects layout (evidence: `src/jax_titan/trainer/train.py`)
- Training Layouts → MPMD Pipeline: builds train step (evidence: `src/jax_titan/trainer/adapters.py`)
- Training Layouts → Layout Sharding: uses sharding (evidence: `src/jax_titan/trainer/adapters.py`)
- MPMD Pipeline → Parallel Shardings: uses shardings (evidence: `src/jax_titan/mpmd/pipeline.py`)
- Training Entry → Stage Checkpoints: saves state (evidence: `src/jax_titan/trainer/train.py`)
- Training Entry → Step Watchdog: monitors steps (evidence: `src/jax_titan/trainer/train.py`)
- Training Layouts → Llama Model: builds stages (evidence: `src/jax_titan/trainer/adapters.py`)
- Training Layouts → MaxText Adapter: adapts model (evidence: `src/jax_titan/trainer/adapters.py`)
- Training Layouts → MaxDiffusion Adapter: adapts model (evidence: `src/jax_titan/trainer/adapters.py`)
- Training Entry → Text Sampling: loads batches (evidence: `src/jax_titan/trainer/train.py`)
- Text Sampling → Llama Tokenizer: tokenizes text (evidence: `src/jax_titan/data/text.py`)
- Llama Model → Semantic Operators: uses operators (evidence: `src/jax_titan/models/llama.py`)
- Semantic Operators → Provider Registry: dispatches ops (evidence: `src/jax_titan/ops/__init__.py`)
- Provider Registry → Native JAX Ops: falls back (evidence: `src/jax_titan/ops/__init__.py`)
- Provider Registry → Accelerator Providers: selects provider (evidence: `src/jax_titan/ops/providers/triton.py`)
- Op Conformance → Semantic Operators: checks implementations (evidence: `src/jax_titan/ops/conformance.py`)
- Training Entry → Run Records: writes events (evidence: `src/jax_titan/trainer/train.py`)
- Training Entry → Environment Capture: captures environment (evidence: `src/jax_titan/trainer/train.py`)
- Training Entry → Power Metrics: collects power (evidence: `src/jax_titan/trainer/train.py`)
- Training Entry → FLOPs Model: reports FLOPs (evidence: `src/jax_titan/trainer/train.py`)
- Run Summaries → Run Records: reads records (evidence: `src/jax_titan/observability/summarize.py`)
- Run Summaries → Power Metrics: reads energy (evidence: `src/jax_titan/observability/summarize.py`)
- Run Summaries → Restart Analysis: computes ETTR (evidence: `src/jax_titan/observability/summarize.py`)
- Training Entry → Profiles and Figures: captures traces (evidence: `src/jax_titan/trainer/train.py`)
- Training User → Budget Planner: requests estimate (evidence: `README.md`)
- Budget Planner → Model Estimates: estimates model (evidence: `src/jax_titan/planning/budget.py`)
- Budget Planner → GPU Catalog: selects GPU (evidence: `src/jax_titan/planning/budget.py`)
- Budget Planner → Memory Planning: estimates memory (evidence: `src/jax_titan/planning/budget.py`)
- Budget Planner → Collective Model: estimates communication (evidence: `src/jax_titan/planning/budget.py`)
- Budget Planner → Step-Time Model: predicts step time (evidence: `README.md`)
- Step-Time Model → Collective Model: models collectives (evidence: `src/jax_titan/planning/step.py`)
- Step-Time Model → Performance Database: checks regressions (evidence: `src/jax_titan/planning/step.py`)