PyTorch Model Logprob Inconsistencies Under High GRPC Traffic

Q: 1. Model Isolation via per‑request clone

Instead of sharing a single nn.Module instance across threads, each request now receives a deep‑copied model wrapped in a torch.no_grad() context. This eliminates shared mutable state.

Q: Operational Experience

During the investigation several misleading clues appeared:

In this article

Table of Contents

Problem Description

During a sustained load test of a Docker‑containerized language model serving endpoint, the logprob values returned by the PyTorch model became inconsistent. Under low traffic the model produced deterministic log‑softmax outputs, but at peak rates (≈200 RPS) the following symptoms were observed:

Log probabilities drifted by up to 0.12 nats between identical inputs.

Intermittent errors appeared in the gRPC logs:

grpc INTERNAL: Received RST_STREAM with error code 2

Container runtime reported OOM events:

docker: Container model_server OOMKilled

PyTorch raised occasional runtime errors:

RuntimeError: CUDA error: device-side assert triggered

When the container restarted after OOM, inference switched from GPU to CPU, changing floating‑point rounding and further degrading logprob precision.

Root Cause Analysis

The inconsistencies stem from a combination of resource contention and thread‑unsafe model usage:

gRPC server thread pool exhaustion – The default max_workers (10) cannot keep up with the incoming request burst. As documented in the gRPC Python concurrency guide, excess requests are queued and eventually processed by newly spawned worker threads, which re‑initialize the model with a fresh random seed, breaking determinism.
Shared model instance without synchronization – PyTorch’s LogSoftmax is not thread‑safe when the same model object is invoked concurrently. Issue pytorch#82456 demonstrates that concurrent calls can corrupt internal buffers, yielding different log‑prob values.
CUDA context race condition – Under high request rates, multiple threads launch kernels on the same CUDA context. Issue pytorch#102341 reports nondeterministic outputs caused by overlapping kernel launches, which can also trigger device‑side asserts.
Docker memory limits and OOM – The container was constrained to 8 GiB (mem_limit: 8g in Compose). Peak memory usage exceeded this limit, causing the Docker daemon to kill the process (OOMKilled). The subsequent restart fell back to CPU inference, altering precision as described in the community post “Docker container OOM leads to fallback to CPU and changes in logprob precision”.
Resource‑constrained CUDA fallback – When a torch.cuda.OutOfMemoryError occurs, PyTorch automatically moves tensors to CPU if torch.backends.cuda.matmul.allow_tf32 is enabled, again changing numerical results.

Collectively, these factors produce the observed logprob drift and error messages.

Investigation and Debugging

The following step‑by‑step investigation helped isolate the failure modes:

Inspect container resource usage

docker stats model_server --no-stream
# Expected output shows CPU ~95 % and MEM 8.2 GiB → OOM risk

Capture gRPC server metrics

grpc_cli ls localhost:50051
# Verify thread pool size and pending request count

Enable PyTorch deterministic mode (does not fix race but reveals nondeterminism)
```
import torch
torch.use_deterministic_algorithms(True)
```

Reproduce race condition locally by invoking the model from multiple threads:

import threading, torch, time
def infer():
    with torch.no_grad():
        out = model(input_tensor)
        print(out.log_softmax(dim=-1).sum().item())
threads = [threading.Thread(target=infer) for _ in range(20)]
[t.start() for t in threads]
[t.join() for t in threads]
# Output shows varying sums → race confirmed

Check Docker logs for OOM events

journalctl -u docker | grep OOMKilled
# Example line:
Jun 15 14:23:01 host docker[1234]: Container model_server OOMKilled

Verify CUDA errors

docker exec -it model_server nvidia-smi
# Shows “GPU 0: XMiB / YMiB” with high utilization
docker logs model_server | grep "CUDA error"
# Example:
RuntimeError: CUDA error: device-side assert triggered

Resolution

Three orthogonal fixes were applied: isolate model execution, enforce resource limits, and tune gRPC concurrency.

1. Model Isolation via per‑request clone

Instead of sharing a single nn.Module instance across threads, each request now receives a deep‑copied model wrapped in a torch.no_grad() context. This eliminates shared mutable state.

Before (single shared model):

# app.py
model = torch.jit.load("model.pt").eval()

def predict(request):
    input_tensor = preprocess(request)
    with torch.no_grad():
        logits = model(input_tensor)
    return postprocess(logits.log_softmax(dim=-1))

After (per‑request clone with lock for GPU reuse):

# app.py
import copy, threading
model_proto = torch.jit.load("model.pt").eval()
model_lock = threading.Lock()   # protects GPU context

def predict(request):
    input_tensor = preprocess(request)
    # Clone weights; share underlying storage to avoid full copy cost
    local_model = copy.deepcopy(model_proto)
    with model_lock:               # ensures one CUDA kernel launch at a time
        with torch.no_grad():
            logits = local_model(input_tensor)
    return postprocess(logits.log_softmax(dim=-1))

2. Configure Docker resource limits and enable swap guard

Increase memory ceiling and disable aggressive OOM killing:

# docker-compose.yml
services:
  model_server:
    image: myorg/model_server:latest
    deploy:
      resources:
        limits:
          memory: 12g
        reservations:
          memory: 8g
    mem_swappiness: 0
    ulimits:
      memlock: -1

Reference: Docker Engine resource constraints.

3. Tune gRPC server concurrency

Set the maximum number of worker threads to match the number of CPU cores and enable a bounded thread‑pool executor to reject excess traffic gracefully.

# server.py
import grpc
from concurrent import futures

MAX_WORKERS = 32  # based on host vCPU count
server = grpc.server(
    futures.ThreadPoolExecutor(max_workers=MAX_WORKERS),
    options=[
        ('grpc.max_concurrent_streams', 1000),
        ('grpc.max_send_message_length', 10 * 1024 * 1024),
    ]
)
# Register services and start

4. Enable CUDA deterministic kernels (optional)

If GPU inference must remain, enforce deterministic behavior:

torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = False

Validation

After deploying the changes, the following verification steps confirmed resolution:

Functional consistency – Run 1 000 identical inference calls in parallel and compare logprob vectors:
```
python validate.py
# Expected output: All vectors equal within 1e-9 tolerance
```
Absence of OOM – Monitor container memory for 30 minutes under 250 RPS:
```
docker stats model_server --no-stream
# MEM % stays < 70 %
```
gRPC health check – Use grpc_health_probe to ensure no RST_STREAM errors:
```
grpc_health_probe -addr=:50051
# Status: SERVING
```

CUDA error log scan – Verify no device‑side asserts:

docker logs model_server | grep "CUDA error"
# No matches

Operational Experience

During the investigation several misleading clues appeared:

LogSoftmax dimension errors were initially blamed on malformed payloads, but the real trigger was the model being invoked before input tensors were correctly reshaped due to race‑condition timing.
The first OOM event occurred only on the GPU‑enabled container; a parallel CPU‑only replica continued serving correctly, highlighting the importance of heterogeneous fall‑back strategies.
Increasing max_workers without a matching increase in CPU quota caused the host to thrash, leading to higher latency even though the logprob values stabilized.

Best Practices and Prevention

Isolate model state per request or use torch.multiprocessing with separate CUDA contexts for true parallelism.
Set explicit Docker memory and CPU limits that exceed peak usage; monitor docker stats and configure mem_swappiness=0 to avoid silent swapping.
Configure gRPC thread pool based on workload characteristics; use back‑pressure (e.g., grpc.max_concurrent_streams) to reject overload instead of queuing indefinitely.
Enable deterministic kernels when reproducibility is required; note the performance trade‑off.
Instrument logs and metrics for:
- Container OOM events (docker events --filter event=oom)
- GPU utilization (nvidia-smi --query-gpu=utilization.gpu,memory.used --format=csv)
- gRPC error counters (e.g., grpc_server_handled_total with label grpc_code)

Related Topic Hub: Distributed Systems Troubleshooting Hub

FAQ

Why does the logprob drift only under high load?
Because concurrent threads share the same model instance, causing race conditions in the LogSoftmax buffer and CUDA context. At low load only one thread runs, so the race never manifests.
Can I keep a single model instance and still be thread‑safe?
Only if you serialize access to the model (e.g., a threading.Lock) or move inference to a separate process pool. Sharing the same CUDA context across threads without synchronization is unsafe.
What is the impact of falling back from GPU to CPU?
CPU inference uses 32‑bit floats and different BLAS kernels, leading to slight rounding differences. More importantly, it reduces throughput, causing downstream latency spikes.
How do I detect that the container has restarted on CPU?
Check the process’s device list at startup:
```
torch.cuda.is_available()
# Returns False after CPU fallback
```
and correlate with Docker’s restart events in the logs.
Is increasing the Docker memory limit sufficient?
It mitigates OOM but does not solve the thread‑safety issue. Both resource limits and proper model isolation are required for deterministic logprob outputs.