feat(subsystembenchmarks): expert-parallel case on OLMoE for checkpointing/pytorch

[yuxin00j]

Stacked on #1088. Please review only the last commit (feat(subsystembenchmarks): expert-parallel case on OLMoE for checkpointing/pytorch); earlier commits belong to #1087 and #1088.

Summary

Adds two OLMoE-1B-7B-0924 rows (38.7 GiB checkpoint, ws8): fsdp and ep dp8ep8. With fsdp, every expert tensor is sliced 8 ways. With ep, each rank owns 8 whole experts per layer.

How it works

• ep keeps the dense params on FSDP2 over dp_shard. keep_local_experts keeps only this rank's experts, keyed by their original index so FQNs stay canonical (model.layers.<L>.mlp.experts.<i>). Those experts are passed as ignored_params to the per-layer fully_shard, so they stay plain tensors that still have optimizer state.
• EP reuses the dp_shard ranks and is not a mesh dimension. validate_case requires expert_parallel_size == data_parallel_size for ep and 1 otherwise.
• This depends on the transformers==4.46.3 pin, where each expert is its own module. Newer versions stack experts into one tensor, and ep would then load like fsdp.
• The OLMoE fsdp row is on the model axis, so run the pair with --sweep-axes="model strategy".

Test Plan

PYTHONPATH=$PWD python -m pytest gcsfs/tests/perf/subsystembenchmarks/checkpointing gcsfs/tests/perf/subsystembenchmarks/tests --run-benchmarks-infra -q
pre-commit run --files <changed files>

106 passed. Same infra tests as the previous PR, plus a tiny OLMoE config (2 layers, 16 experts):

• test_parallelize.py: each of 8 fake ranks keeps exactly its 2 experts per layer as plain trainable tensors, dense params stay FSDP Shard(0), and the expert FQN sets across ranks are disjoint and cover every expert.
• test_roundtrip.py: a real gloo ep save/load round trip matches checksums, and the driver reports fewer read items per rank for ep than for fsdp on the same model.
• test_configs.py: the two OLMoE case names and the expert_parallel_size validation.

This PR was written with help from an AI coding assistant.

[gemini-code-assist]

Code Review

This pull request introduces a new native PyTorch checkpointing benchmark group (checkpointing/pytorch) to measure torch.distributed.checkpoint.load performance from GCS across various parallelism strategies (FSDP2, HSDP, TP, PP, EP). It also refactors shared distributed utilities into a common _dist.py module to eliminate duplication across existing benchmark groups, updates the BigQuery schema to support the new metrics, and adds comprehensive tests. Feedback on the changes identifies a critical bug in the custom load planner where accessing itemsize on a torch.dtype will raise an AttributeError at runtime; using torch._utils._element_size is recommended instead.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 90.76%. Comparing base (1bc5128) to head (654632e).

Additional details and impacted files

@@           Coverage Diff           @@
##             main    #1091   +/-   ##
=======================================
  Coverage   90.76%   90.76%           
=======================================
  Files          16       16           
  Lines        3650     3650           
=======================================
  Hits         3313     3313           
  Misses        337      337           

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.