Checkpoint write throughput is capped by CPU-side data generation, not storage (`dgen-py` under-threaded → Generation is the bottleneck)

[wolfgang-desalvador]

Severity: High — the reported checkpoint metric measures synthetic-data
generation (CPU) throughput rather than storage throughput, so it under-reports
the storage system and materially changes results.

Summary

During a llama3-70b checkpoint save run, the streaming-checkpoint
producer (synthetic tensor generation via dgen-py) is a co-equal — and in
a share of chunks the dominant — bottleneck relative to storage I/O. The
producer is CPU-bound and under-threaded, so the reported checkpoint throughput
is limited by how fast the client can generate the checkpoint tensors, not by
how fast the storage target can absorb them.

Per-chunk profiling from a single run:

Signal	Observed
Generator threads per rank	3 threads (threads=3)
Throughput ratio (gen/io)	~0.9–1.1×
Chunks where Generation is the bottleneck	~35% (1,327 of 3,840)
Pipeline overhead when Generation-bound	up to ~10% (vs ~3% when I/O-bound)

Because generation and I/O run at a similar rate, the measured throughput sits
at whichever side is momentarily slower — and about a third of the time that is
the CPU-side generator, not the storage device.

Environment

• Model / mode: llama3-70b, checkpointing, num_checkpoints_write=10
• Ranks: num_processes=64 (16 × 4)
• Client nodes: 16 nodes, each with 192 vCPUs and 1.8 TiB of RAM, reserved
  entirely for the job (all cores available, ranks launched unbound — the
  3-thread generator is an application choice, not a scheduler/cgroup CPU limit)
• Target: POSIX filesystem — backend: file, o_direct: false
• mlpstorage: 6576710 (3.0.26)
• Per-shard checkpoint size: ~12.2 GB, streamed through a 4 × 32 MB buffer pool

Analysis

The streaming checkpoint uses a producer–consumer pipeline: a producer
generates synthetic tensor data (dgen-py) into a small shared-memory buffer
pool, and a writer subprocess drains those buffers to storage. The intent is
that the writer (storage I/O) is the bottleneck so the run measures storage
throughput.

In practice the producer is initialised with only 3 worker threads
(threads=3), so its generation rate is close to the storage
write rate. The consequences:

• The pipeline is generation-starved a share of the time. In ~35% of
  profiled chunks the profiler labels Bottleneck: Generation.
• The reported metric is not a pure storage measurement. When generation is
  the limiter, the reported throughput reflects CPU tensor generation, not the
  storage target — which can absorb more than the producer supplies.
• Faster storage cannot be distinguished from slower storage at this
  configuration, because both are bounded by the generation ceiling.

Profiling evidence

A larger number of occurrences where the ratio drops below 1.0× — here storage I/O
outruns generation, and with pipeline overhead at only 1.7% the two
stages are well overlapped, so the storage target could sustain more than the
producer delivers:

================================================================================
RESULTS
================================================================================
Generation: xxxx @ xxxx GB/s
I/O:         xxxx @ xxxx GB/s
  - write:   xxxx
  - close:   xxxx (fsync/finalize)
Total:       xxxx

Throughput ratio: 0.9x (gen/io)
Pipeline overhead: 1.7%
Bottleneck: Generation
Chunks: 391
================================================================================

Across the full run the profiler reports Bottleneck: Generation on 1,327
chunks and Bottleneck: I/O on 2,513 chunks (3,840 profiled chunks total) —
i.e. generation is the limiting stage about a third of the time, and the
gen/io ratio stays close to 1.0×, so the storage target is not allowed to run
ahead of the producer.

Impact (why this affects results)

The checkpointing benchmark is intended to measure storage save/load
throughput. With the producer running 3 worker threads, the synthetic-data
generator runs at a rate close to the storage writer, so the reported metric is
co-determined (and at times gated) by client CPU throughput rather than the
storage system under test. A comparison of storage targets at this
configuration is confounded: a faster storage system can produce the same
reported number because the generator, not the device, sets the ceiling.

[wolfgang-desalvador]

Tracing the under-threaded generator (threads=3/192) to the dgen-py throttle in mlpstorage_py/checkpointing/streaming_checkpoint.py, in _init_generator():

mpi_world_size = 1
for _env_var in ('OMPI_COMM_WORLD_SIZE', 'PMI_SIZE', 'MV2_COMM_WORLD_SIZE'):
    _v = os.environ.get(_env_var)
    if _v:
        mpi_world_size = max(1, int(_v))
        break
total_cpus = os.cpu_count() or 4
max_threads = max(1, total_cpus // mpi_world_size)

The throttle divides the per-node CPU count (os.cpu_count(), which is node-local) by OMPI_COMM_WORLD_SIZE — the total number of ranks across the entire job, not the number of ranks on this node. The intent (per the inline comment) was to avoid oversubscribing a host when several ranks share it, but the divisor used is the global world size, so the thread budget shrinks as more nodes are added even though per-node contention is unchanged.

Why it collapses with higher number of CPUs

Both of the jobs I compared run 4 ranks per node on 192-vCPU hosts (-host host:4,...):

Job	Nodes	Ranks/node	OMPI_COMM_WORLD_SIZE	192 // world_size (actual)	192 // ranks_per_node (correct)
A	2	4	8	24	48
B	16	4	64	3	48

Per-node contention is identical in both runs (4 ranks sharing 192 cores), so the correct thread budget is 192 / 4 = 48 in both cases. Because the divisor is the global world size, the 16-node job gets 192 // 64 = 3 threads/rank — a 16× under-allocation — while the hardware per node is unchanged.

[russfellows]

OK, as shown above, I am 99% confident that my fixes in PR #702 will address these issues. I have recreated the bug without the patches, then verified that with the patches the number of CPU cores allocated to dgen-py is as intended, which is 50% of what is available.

[FileSystemGuy]

@wolfgang-desalvador When you get a chance to test this, please either comment here with new info or close this as fixed.

[russfellows]

@wolfgang-desalvador , can you recheck?

I am pretty confident that PR's #700 and #702 SHOULD have fixed this issue. I verified and reproduced your issue, and verified that together PRs 700 and 702 fix the problem.