calc_throughput.py

import time
import torch
import torch.nn as nn
from torch.utils.data import DataLoader, Dataset
from transformers import AutoModelForCausalLM
from torch.optim import AdamW

# --- Configuration for Qwen3-1.7B (Proxy) ---
# Based on the specs gathered: 1.7B params, 28 layers, 16 heads, 2048 hidden size.

# Training Hyperparameters
BATCH_SIZE = 4        # Adjust based on memory (H200 141GB is huge, can go higher)
SEQ_LENGTH = 4096     # Sequence length for benchmarking
GRAD_ACCUMULATION = 1
TOTAL_STEPS = 50      # Number of measurement steps
WARMUP_STEPS = 10     # Steps to ignore for stable timing

class SyntheticDataset(Dataset):
    """Generates random tokens to remove disk I/O bottlenecks."""
    def __init__(self, seq_len, vocab_size, length=1000):
        self.seq_len = seq_len
        self.vocab_size = vocab_size
        self.length = length

    def __len__(self):
        return self.length

    def __getitem__(self, idx):
        return {
            "input_ids": torch.randint(0, self.vocab_size, (self.seq_len,), dtype=torch.long),
            "labels": torch.randint(0, self.vocab_size, (self.seq_len,), dtype=torch.long)
        }

def train_benchmark():
    # 1. Setup Device
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"Using device: {torch.cuda.get_device_name(0)}")
    
    # 2. Initialize Model (Random Weights)
    print("Initializing Qwen3-1.7B configuration...")
    model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen3-1.7B", torch_dtype="auto").cuda()#, device_map="auto")#, attn_implementation="flash_attention_2")
    model.train()
    
    # Enable Gradient Checkpointing (optional, saves memory, costs ~20% compute)
    # model.gradient_checkpointing_enable() 

    optimizer = AdamW(model.parameters(), lr=1e-4)

    # 3. Setup Data
    dataset = SyntheticDataset(SEQ_LENGTH, model.config.vocab_size)
    dataloader = DataLoader(dataset, batch_size=BATCH_SIZE, num_workers=0)
    
    print(f"\nStarting Benchmark:")
    print(f"Batch Size: {BATCH_SIZE} | Seq Len: {SEQ_LENGTH}")
    print("-" * 50)

    # 4. Training Loop
    step_times = []
    tokens_processed = 0
    start_time = 0

    iter_loader = iter(dataloader)
    
    for step in range(TOTAL_STEPS + WARMUP_STEPS):
        try:
            batch = next(iter_loader)
        except StopIteration:
            iter_loader = iter(dataloader)
            batch = next(iter_loader)

        input_ids = batch["input_ids"].to(device)
        labels = batch["labels"].to(device)

        # Synchronize before timing
        torch.cuda.synchronize()
        t0 = time.time()

        # --- Forward Pass ---
        outputs = model(input_ids=input_ids, labels=labels)
        loss = outputs.loss

        # --- Backward Pass ---
        loss.backward()

        # --- Optimizer Step ---
        optimizer.step()
        optimizer.zero_grad()

        # Synchronize after timing
        torch.cuda.synchronize()
        t1 = time.time()
        
        dt = t1 - t0

        if step >= WARMUP_STEPS:
            step_times.append(dt)
            current_tokens = BATCH_SIZE * SEQ_LENGTH
            tokens_processed += current_tokens
            
            # Real-time logging
            tokens_per_sec = current_tokens / dt
            print(f"Step {step+1}/{TOTAL_STEPS+WARMUP_STEPS} | "
                  f"Time: {dt:.3f}s | "
                  f"Throughput: {tokens_per_sec:,.0f} tok/s")
        else:
            print(f"Step {step+1} (Warmup)...")

    # 5. Final Statistics
    avg_time = sum(step_times) / len(step_times)
    avg_throughput = (BATCH_SIZE * SEQ_LENGTH) / avg_time

    print("\n" + "=" * 50)
    print("BENCHMARK RESULTS")
    print("=" * 50)
    print(f"Average Step Time:   {avg_time:.4f} seconds")
    print(f"Average Throughput:  {avg_throughput:,.2f} tokens/second")
    print("=" * 50)

if __name__ == "__main__":
    train_benchmark()