docs/apps/zen-docs/content/blog/training-gym.mdx

---
title: "Training Gym: A Platform for Open Model Development"
description: "Announcing Training Gym, our open platform for collaborative large model training."
date: "2023-09-11"
author: "Zach Kelling"
tags: ["Infrastructure", "Training", "Open Source"]
---

Training large language models requires more than algorithms. It requires infrastructure: distributed training frameworks, data pipelines, experiment tracking, and evaluation harnesses. Today we open source Training Gym, our complete platform for model development.

## Why Training Gym?

Open AI development faces an infrastructure gap. Publishing model weights is valuable, but it’s not enough. Researchers need:

  * Reproducible training pipelines
  * Scalable distributed training
  * Standardized evaluation
  * Experiment management
  * Data processing tools

Training Gym provides all of this in an integrated, open source package.

## Architecture
```

    +------------------+     +------------------+     +------------------+
    |   Data Pipeline  | --> |  Training Loop   | --> |   Evaluation     |
    +------------------+     +------------------+     +------------------+
             |                       |                        |
             v                       v                        v
    +------------------+     +------------------+     +------------------+
    |   Data Registry  |     | Checkpoint Store |     |   Metrics DB     |
    +------------------+     +------------------+     +------------------+
                                    |
                                    v
                        +------------------+
                        |   Experiment     |
                        |   Tracker        |
                        +------------------+
    

```

### Data Pipeline

The data pipeline handles:

  * **Ingestion** : Load from local files, cloud storage, or streaming sources
  * **Processing** : Tokenization, filtering, deduplication
  * **Mixing** : Combine multiple data sources with configurable ratios
  * **Streaming** : Memory-efficient data loading for large corpora

```python

    from training_gym.data import DataPipeline, MixedDataset
    
    pipeline = DataPipeline(
        sources=[
            ("s3://data/books", 0.3),
            ("s3://data/web", 0.5),
            ("s3://data/code", 0.2),
        ],
        tokenizer="zoo-labs/zen-tokenizer",
        sequence_length=2048,
    )
    
    dataset = pipeline.build()
    

```

### Distributed Training

Training Gym supports multiple distributed training strategies:

  * **Data Parallel** : Simple replication across devices
  * **Tensor Parallel** : Split layers across devices
  * **Pipeline Parallel** : Split model stages across devices
  * **ZeRO** : Memory-efficient data parallelism
  * **FSDP** : Fully sharded data parallel (PyTorch native)

Configuration is declarative:
```

    distributed:
      strategy: fsdp
      world_size: 64
      sharding_strategy: full_shard
      mixed_precision: bf16
      gradient_checkpointing: true
    

```

### Training Loop

The training loop is modular and extensible:
```python

    from training_gym import Trainer, TrainingConfig
    
    config = TrainingConfig(
        model="zen-7b",
        optimizer="adamw",
        learning_rate=1e-4,
        batch_size=2048,
        max_steps=100000,
        warmup_steps=2000,
        weight_decay=0.1,
    )
    
    trainer = Trainer(config)
    trainer.fit(dataset)
    

```

Built-in features:

  * Learning rate scheduling
  * Gradient clipping
  * Mixed precision training
  * Automatic checkpointing
  * Loss spike detection and recovery

### Evaluation

Standardized evaluation across common benchmarks:
```python

    from training_gym.eval import Evaluator
    
    evaluator = Evaluator(
        benchmarks=["mmlu", "hellaswag", "winogrande", "arc"],
        model=model,
    )
    
    results = evaluator.run()
    

```

Supported benchmarks:

  * MMLU (multitask language understanding)
  * HellaSwag (commonsense reasoning)
  * WinoGrande (coreference resolution)
  * ARC (science questions)
  * TruthfulQA (truthfulness)
  * HumanEval (code generation)
  * GSM8K (math reasoning)

### Experiment Tracking

Every training run is tracked:
```python

    from training_gym import Experiment
    
    with Experiment("zen-7b-v2") as exp:
        exp.log_config(config)
        trainer.fit(dataset)
        exp.log_metrics(results)
        exp.log_artifacts(["model.pt", "tokenizer/"])
    

```

The experiment tracker records:

  * Hyperparameters
  * Training metrics (loss, gradient norms, learning rates)
  * Evaluation results
  * System metrics (GPU utilization, memory)
  * Artifacts (checkpoints, configs)

## Reproducibility

Training Gym emphasizes reproducibility:

### Deterministic Training
```

    reproducibility:
      seed: 42
      deterministic_algorithms: true
      cublas_workspace_config: ":4096:8"
    

```

Same seed, same results (within floating point precision).

### Environment Capture

Every experiment records:

  * Git commit hash
  * Package versions
  * Hardware configuration
  * CUDA/cuDNN versions

### Configuration as Code

All configs are versioned YAML:
```

    # experiments/zen-7b-v2.yaml
    model:
      architecture: llama
      hidden_size: 4096
      num_layers: 32
      num_heads: 32
      vocab_size: 32000
    
    training:
      batch_size: 2048
      learning_rate: 3e-4
      max_steps: 150000
    

```

## Community Features

Training Gym includes tools for collaborative development:

### Model Registry

Share and discover models:
```python

    from training_gym.registry import ModelRegistry
    
    registry = ModelRegistry()
    
    # Publish a model
    registry.push("my-org/my-model", model, config)
    
    # Load a model
    model = registry.pull("zoo-labs/zen-7b")
    

```

### Leaderboards

Automatic benchmark submission:
```

    evaluator.submit_to_leaderboard(
        model_name="zen-7b-v2",
        organization="zoo-labs",
    )
    

```

### Dataset Sharing
```python

    from training_gym.data import DatasetRegistry
    
    # Share processed datasets
    DatasetRegistry.push("my-corpus", dataset, license="cc-by-4.0")
    
    # Load shared datasets
    dataset = DatasetRegistry.pull("zoo-labs/zen-pretrain-v1")
    

```

## Getting Started

### Installation
```bash

    pip install training-gym
    
    # For distributed training
    pip install training-gym[distributed]
    
    # For evaluation suite
    pip install training-gym[eval]
    

```

### Quick Start
```python

    from training_gym import quickstart
    
    # Train a small model to verify setup
    quickstart.train_tiny_model()
    
    # Run evaluation suite
    quickstart.evaluate_model("my-model")
    

```

### Documentation

Full documentation at docs.training-gym.ai:

  * Getting started guide
  * Architecture overview
  * API reference
  * Example configurations
  * Troubleshooting

## Roadmap

**Q4 2023** : Multi-modal training support **Q1 2024** : Reinforcement learning from human feedback (RLHF) integration **Q2 2024** : Federated training support **Q3 2024** : Automated hyperparameter optimization

## Conclusion

Open AI development needs open infrastructure. Training Gym provides the tools to train, evaluate, and share models. Join us in building the future of open AI.

Repository: github.com/zoo-labs/training-gym

* * *

 _Zach Kelling is a co-founder of Zoo Labs Foundation._