Parallax: MaxText vs PyTorch on NVIDIA Blackwell (RTX PRO 6000 96GB)

Parallax is a production-grade benchmarking suite and PyPI package designed to compare Google MaxText (JAX/XLA) and PyTorch (TransformerEngine / Triton) for pre-training a ~100M parameter LLaMA transformer model on the TinyStories dataset on NVIDIA Blackwell architecture (RTX PRO 6000 Blackwell 96GB, SM100).

It evaluates:

• Training Throughput & Speed: Measured tokens/second, step latency (ms), achieved dense TFLOPs, and Model FLOPs Utilization (MFU %).
• Loss & Convergence: Step-by-step cross-entropy loss trajectory, final perplexity, and quantization error degradation.
• Memory Footprint: Peak VRAM allocated and reserved across precision formats.
• Precision Formats:
  1. bf16: Bfloat16 standard baseline mixed precision.
  2. fp8: 8-bit floating point (E4M3 forward, E5M2 backward) using TransformerEngine DelayedScaling (PyTorch) and fp8_gpu (MaxText).
  3. nvfp4: Native NVIDIA Blackwell 4-bit floating point (E2M1) with 16-element block scaling and Random Hadamard Transform (RHT) via TransformerEngine NVFP4BlockScaling (PyTorch) and te_nvfp4 (MaxText).

---

1. Model Architecture (~100.3M Parameters)

To ensure strict parity, both PyTorch and MaxText instantiate an identical LLaMA decoder-only transformer:

Hyperparameter	Value	Notes
Parameters	100,313,856 (~100.3M)	Untied embeddings, vocab_size=10,000
Decoder Layers	12	Full pre-norm transformer blocks
Hidden Dimension ($d$)	768	Width of representations
Attention Heads ($H$)	12	Multi-Head Attention (MHA)
Head Dimension ($d_k$)	64	$768 / 12 = 64$
MLP Intermediate Dim	2048	SwiGLU: down(silu(gate) * up)
Sequence Length	1024	Configurable via CLI (--seq-len)
Positional Encoding	RoPE	Rotary Position Embeddings ($\theta = 10000.0$)
Normalization	RMSNorm	Root Mean Square Norm ($\epsilon = 10^{-5}$)
Dataset	TinyStories	roneneldan/TinyStories (or synthetic zero-overhead mode)

---

2. Precision & Blackwell Quantization Recipes

┌─────────────────────────────────────────────────────────────────────────────┐
│                    NVIDIA BLACKWELL SM100 TENSOR CORES                      │
├───────────────────┬────────────────────────────┬────────────────────────────┤
│ Precision         │ PyTorch Stack              │ MaxText (JAX/XLA) Stack    │
├───────────────────┼────────────────────────────┼────────────────────────────┤
│ BF16              │ torch.amp.autocast(bf16)   │ dtype=bfloat16             │
│ FP8 (Hybrid)      │ TE DelayedScaling (E4M3/5) │ quantization="fp8_gpu"     │
│ NVFP4 (Native 4b) │ TE NVFP4BlockScaling + RHT │ quantization="te_nvfp4"    │
└───────────────────┴────────────────────────────┴────────────────────────────┘

BF16

• Standard baseline mixed-precision on Blackwell.
• Optimizer states stored in FP32/BF16.

FP8 (Format.HYBRID)

• Forward pass GEMMs evaluated in E4M3 (higher precision).
• Backward pass gradient GEMMs evaluated in E5M2 (wider dynamic range).
• Dynamic scaling computed via amax sliding history window.

NVFP4 (NVIDIA Blackwell 4-bit Float)

• E2M1 Microscaling: 1 sign bit, 2 exponent bits, 1 mantissa bit.
• Hierarchical Block Scaling: 16-element vector blocks with FP8 scale factors.
• Random Hadamard Transform (RHT): Rotates activation space before quantization to eliminate outlier spikes and preserve convergence.

---

3. Quickstart & Installation

Option A: Install from PyPI

pip install slm-parallax

Option B: Install from source

git clone https://github.com/your-org/slm-parallax.git
cd slm-parallax
pip install -e .

Option C: Remote Blackwell Setup Script

If you are running directly on a remote server with an NVIDIA RTX PRO 6000 Blackwell (96GB):

bash scripts/setup_remote_blackwell.sh

---

4. CLI Usage

Parallax installs a unified console script: parallax.

1. Audit System & GPU Environment

parallax check-env

Outputs CUDA version, Blackwell architecture detection (SM100+), VRAM capacity, PyTorch, TransformerEngine, and JAX status.

2. Run the Benchmark Matrix

Compare PyTorch vs MaxText across bf16, fp8, and nvfp4:

parallax benchmark \
    --frameworks torch,maxtext \
    --precisions bf16,fp8,nvfp4 \
    --batch-size 32 \
    --steps 200 \
    --warmup 20 \
    --synthetic \
    --output-dir benchmark_results

Flags:

• --frameworks: Comma-separated (torch, maxtext).
• --precisions: Comma-separated (bf16, fp8, nvfp4).
• --batch-size: Micro-batch size per device (default: 32).
• --seq-len: Sequence length (default: 1024).
• --steps: Number of measured steps (default: 200).
• --warmup: Warmup steps excluded from throughput calculations (default: 20).
• --synthetic / --hf: Use zero-overhead synthetic batches for pure compute benchmarking, or stream HuggingFace TinyStories.
• --output-dir: Results destination directory.
• --report: Generates GitHub-flavored Markdown report (report.md).

3. Run a Single Training Job

# PyTorch with NVFP4 on Blackwell
parallax train --framework torch --precision nvfp4 --batch-size 32 --steps 200

# MaxText with FP8
parallax train --framework maxtext --precision fp8 --batch-size 32 --steps 200

4. Compare Existing Results

parallax compare --results-dir benchmark_results --output-report report.md

5. Export Standalone MaxText YAML Configs

parallax export-configs --output-dir maxtext_configs

---

5. Directory Structure

parallax/
├── parallax/
│   ├── cli.py                     # Click CLI interface (benchmark, train, compare, check-env)
│   ├── config.py                  # Llama100MConfig, TrainingConfig, BenchmarkConfig
│   ├── models/
│   │   ├── llama.py               # LLaMA-100M PyTorch model with TransformerEngine Linear
│   │   └── rope.py                # Complex RoPE rotary position embeddings
│   ├── data/
│   │   └── tinystories.py         # TinyStories dataset loader, tokenizer, synthetic generator
│   ├── engines/
│   │   ├── base.py                # MetricStep, BenchmarkResult dataclasses
│   │   ├── torch_engine.py        # PyTorch training loop (CUDA event timing, AdamW, cosine)
│   │   ├── maxtext_engine.py      # MaxText orchestrator, YAML builder, telemetry parser
│   │   └── precision_manager.py   # TransformerEngine recipes & MaxText flags
│   ├── benchmark/
│   │   ├── runner.py              # Benchmark matrix runner
│   │   ├── comparator.py          # Statistical speedup & loss delta analysis
│   │   └── reporting.py           # Terminal tables & Markdown report generation
│   ├── configs/                   # MaxText YAML config templates (BF16, FP8, NVFP4)
│   └── utils/
│       ├── hardware.py            # RTX PRO 6000 Blackwell specs & SM100 detection
│       └── flops.py               # MFU % and theoretical FLOPs calculator
├── scripts/
│   ├── setup_remote_blackwell.sh  # Automated setup for remote RTX PRO 6000 server
│   └── run_benchmarks.sh          # One-click benchmark runner
├── tests/                         # Full unit test suite (models, precisions, data, CLI)
└── pyproject.toml                 # PEP 517/621 package metadata

---

6. Remote Blackwell RTX PRO 6000 Execution Guide

Follow these steps to run on the remote machine:

1. SSH into the remote server:

   ssh user@remote-blackwell-ip
   

2. Transfer / Clone the repository:

   git clone <repo_url> parallax
   cd parallax
   

3. Execute the setup script:

   bash scripts/setup_remote_blackwell.sh
   

4. Activate the environment & verify Blackwell SM100:

   source ~/parallax_venv/bin/activate
   parallax check-env
   

5. Run the full benchmark suite:

   bash scripts/run_benchmarks.sh benchmark_results 32 200 20
   

6. Review the generated report:

   cat benchmark_results/report.md
   

---

7. Metrics & Telemetry Explained

• Throughput (tokens/sec): Total tokens processed per unit time: $$\text{Throughput} = \frac{\text{Batch Size} \times \text{Sequence Length}}{\Delta t_{\text{step}}}$$
• Model FLOPs Utilization (MFU %): Hardware compute efficiency against theoretical peak Tensor Core dense TFLOPs: $$\text{FLOPs/token} = 6 \cdot P + 12 \cdot L \cdot S \cdot d$$ $$\text{MFU} = \frac{\text{Throughput} \times \text{FLOPs/token}}{\text{Peak TFLOPs}} \times 100%$$
• Convergence Delta ($\Delta \text{Loss}$): Loss penalty incurred by low-precision representation: $$\Delta \text{Loss}{\text{precision}} = \text{Loss}{\text{precision}} - \text{Loss}_{\text{BF16}}$$

---

License

Apache License 2.0. See LICENSE for details.