Metadata-Version: 2.4
Name: slm-parallax
Version: 0.1.1
Summary: Benchmarking MaxText vs PyTorch on NVIDIA Blackwell (RTX PRO 6000 96GB) across BF16, FP8, and NVFP4
Author: Parallax Team
License: Apache-2.0
Project-URL: Homepage, https://github.com/your-org/slm-parallax
Project-URL: Repository, https://github.com/your-org/slm-parallax.git
Project-URL: Documentation, https://github.com/your-org/slm-parallax#readme
Project-URL: Issues, https://github.com/your-org/slm-parallax/issues
Keywords: deep-learning,blackwell,transformer,nvfp4,fp8,bf16,maxtext,pytorch,benchmark,tinystories,llama
Classifier: Development Status :: 4 - Beta
Classifier: Intended Audience :: Science/Research
Classifier: Intended Audience :: Developers
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3.10
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Classifier: Programming Language :: Python :: 3.14
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Requires-Python: >=3.10
Description-Content-Type: text/markdown
License-File: LICENSE
Requires-Dist: torch>=2.2.0
Requires-Dist: click>=8.0.0
Requires-Dist: pyyaml>=6.0
Requires-Dist: numpy>=1.24.0
Requires-Dist: tabulate>=0.9.0
Requires-Dist: transformers>=4.36.0
Requires-Dist: datasets>=2.16.0
Provides-Extra: te
Requires-Dist: transformer-engine>=2.7.0; extra == "te"
Provides-Extra: jax
Requires-Dist: jax>=0.4.30; extra == "jax"
Requires-Dist: jaxlib; extra == "jax"
Provides-Extra: all
Requires-Dist: transformer-engine>=2.7.0; extra == "all"
Requires-Dist: jax>=0.4.30; extra == "all"
Requires-Dist: jaxlib; extra == "all"
Dynamic: license-file

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

[![License: Apache 2.0](https://img.shields.io/badge/License-Apache%202.0-blue.svg)](LICENSE)
[![Python 3.10+](https://img.shields.io/badge/python-3.10+-blue.svg)](https://www.python.org/downloads/)
[![Hardware](https://img.shields.io/badge/Target%20GPU-RTX%20PRO%206000%20Blackwell%20(96GB)-green.svg)]()
[![Precision](https://img.shields.io/badge/Precision-BF16%20%7C%20FP8%20%7C%20NVFP4-purple.svg)]()

**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
```bash
pip install slm-parallax
```

### Option B: Install from source
```bash
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
bash scripts/setup_remote_blackwell.sh
```

---

## 4. CLI Usage

Parallax installs a unified console script: `parallax`.

### 1. Audit System & GPU Environment
```bash
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`:
```bash
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
```bash
# 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
```bash
parallax compare --results-dir benchmark_results --output-report report.md
```

### 5. Export Standalone MaxText YAML Configs
```bash
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**:
   ```bash
   ssh user@remote-blackwell-ip
   ```
2. **Transfer / Clone the repository**:
   ```bash
   git clone <repo_url> parallax
   cd parallax
   ```
3. **Execute the setup script**:
   ```bash
   bash scripts/setup_remote_blackwell.sh
   ```
4. **Activate the environment & verify Blackwell SM100**:
   ```bash
   source ~/parallax_venv/bin/activate
   parallax check-env
   ```
5. **Run the full benchmark suite**:
   ```bash
   bash scripts/run_benchmarks.sh benchmark_results 32 200 20
   ```
6. **Review the generated report**:
   ```bash
   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](LICENSE) for details.
