Echo TTS — Scaled Training

DiT-based diffusion transformer for text-to-speech, conditioned on frozen T5Gemma2 text encoders and learned speaker embeddings from DAC-VAE latents.

Two training backends share the same model architecture:

• TorchTitan — DDP / FSDP2 with optional torch.compile
• Megatron — Tensor parallelism (TP) via Megatron-Core + TransformerEngine

Repository structure

model/
  config.py          Model config dataclass + 800M/3B/8B scale presets
  components.py      Shared layers (RMSNorm, QKNorm, RoPE, AdaLN, MLP)
  dit.py             Standard PyTorch model (DDP / FSDP2 / compile)
  dit_megatron.py    Megatron variant (tensor parallelism + TE kernels)

data/
  dataset.py         LocalTarDataset — streams from tar archives
  tokenizer.py       Extended tokenizer with conditioning tokens

encoder/
  frozen.py          FrozenTextEncoder wrapper (T5 / T5Gemma2)

train/
  torchtitan.py      DDP/FSDP2 training entry point
  megatron.py        Megatron TP+DDP training entry point
  utils.py           MFU calculation, distributed helpers

configs/             YAML configs per scale and framework
scripts/             SLURM launcher

Setup

git clone git@github.com:GLJS/scaled_echo_tts.git
cd scaled_echo_tts
pip install -e .

# For Megatron support:
pip install -e ".[megatron]"

Quick start (single machine, no SLURM)

TorchTitan — DDP on 4 GPUs

torchrun --nproc_per_node=4 \
    -m train.torchtitan --config configs/tt_800M.yaml

TorchTitan — FSDP2 + compile on 4 GPUs

torchrun --nproc_per_node=4 \
    -m train.torchtitan --config configs/tt_3B.yaml

Megatron — TP=2 on 4 GPUs

torchrun --nproc_per_node=4 \
    -m train.megatron --config configs/mg_3B.yaml

Megatron — TP=4 on 4 GPUs

torchrun --nproc_per_node=4 \
    -m train.megatron --config configs/mg_8B.yaml

Multi-node with SLURM

The launcher script auto-detects the framework from the YAML config:

# 1 node (4 GPUs)
sbatch --nodes=1 scripts/run.sh configs/tt_3B.yaml

# 4 nodes (16 GPUs), 2 hour wall time
sbatch --nodes=4 --time=02:00:00 scripts/run.sh configs/mg_8B.yaml

# 16 nodes (64 GPUs)
sbatch --nodes=16 --time=02:00:00 scripts/run.sh configs/tt_800M.yaml

Multi-node without SLURM

Set the distributed environment variables manually:

# On each node:
export MASTER_ADDR=<first-node-hostname>
export MASTER_PORT=29500
export NCCL_IB_HCA=mlx5            # InfiniBand adapter
export NCCL_SOCKET_IFNAME=ib0      # IB network interface

# Node 0:
torchrun --nproc_per_node=4 --nnodes=2 --node_rank=0 \
    --master_addr=$MASTER_ADDR --master_port=$MASTER_PORT \
    -m train.torchtitan --config configs/tt_3B.yaml

# Node 1:
torchrun --nproc_per_node=4 --nnodes=2 --node_rank=1 \
    --master_addr=$MASTER_ADDR --master_port=$MASTER_PORT \
    -m train.torchtitan --config configs/tt_3B.yaml

Configuration

All settings live in a single YAML file. Copy and modify:

cp configs/tt_3B.yaml configs/my_experiment.yaml

Key config fields

Data format

Training data is stored as tar archives containing triplets:

sample_001.npy       # Audio latent (T, 128) float16 — DAC-VAE encoded
sample_001.ref.npy   # Speaker reference latent (T, 128) float16
sample_001.json      # {"text": "...", "target_duration": 5.2, ...}

Set the data directory in the training script or config. The LocalTarDataset streams from tars with threaded prefetching.

Model architecture

Text ──→ [Frozen T5Gemma2] ──→ Dense Projection ──→ text_embeddings
                                                          │
Audio ──→ [DAC-VAE latent] ──→ noise + timestep t         │
                │                                          │
Speaker ──→ [Patch + 8L Transformer] ──→ speaker_emb       │
                                              │            │
                                    ┌─────────┴────────────┘
                                    ▼
                            DiT Decoder (24–32 layers)
                            - AdaLN timestep conditioning
                            - Joint self + cross attention
                              (text stream + speaker stream)
                            - QKNorm + RoPE
                            - Gated SiLU MLP
                                    │
                                    ▼
                            noise prediction (T, 128)

Scales:

Swapping components

Different text encoder

Change encoder_model in the YAML config:

encoder_model: google/t5gemma-2-4b-4b   # larger encoder
# or
encoder_model: google-t5/t5-base         # standard T5

Update t5_hidden_dim to match the encoder's output dimension.

Unfrozen encoder

frozen_encoder: false

The encoder gradients will flow and the encoder parameters will be included in the optimizer.

Different parallelism

# DDP (replicate model, shard data)
parallelism: ddp

# FSDP2 (shard model + optimizer across GPUs)
parallelism: fsdp2

# Megatron TP=2 (split attention heads across 2 GPUs, DDP across rest)
framework: megatron
tensor_parallel_size: 2

Results

Benchmark results on NVIDIA GH200 (990 TFLOPS bf16 peak):

Full results dashboard: https://share.gijs.me/model_flops_utilization.html