Skip to content

Commit 3f1b6d8

Browse files
lolpackmeta-codesync[bot]
lolpack
authored and
meta-codesync[bot]
committed
Add pt2_compile_benchmark benchmark for cinderX
Summary: # CinderX + TorchRec pt2_compile_benchmark (CPU) Benchmarks CinderX JIT impact on the TorchRec `pt2_compile_benchmark` (eager forward+backward of DistributedModelParallel), adapted to run on CPU. Uses `FakeStore` and a fake process group — no GPU or multi-node setup needed. ## Files | File | Description | |---|---| | `run_with_cinderx.py` | CinderX JIT benchmark runner | | `test_pt2_multiprocess.py` | Adapted benchmark (CPU, baseline model/sharders) | | `BUCK` | Build targets | | `README.md` | This file | ## CPU Adaptations The original benchmark (`fbcode/torchrec/distributed/tests/test_pt2_multiprocess.py`) runs on GPU. This version is adapted to run on CPU so it can run anywhere without GPU or multi-node setup: 1. `torch.device("cuda")` -> `torch.device("cpu")` 2. `compute_device="cuda"` -> `compute_device="cpu"` in Topology 3. `fullgraph=True` -> `fullgraph=False` (fbgemm CPU ops don't support full graph) 4. `convert_to_vb=True` -> `convert_to_vb=False` (VBE path uses unsupported CPU ops) 5. Relaxed eager-vs-compile assertion (CPU fake PG produces NaN in compile path) Reviewed By: stroxler, DinoV Differential Revision: D95859076 fbshipit-source-id: dd893e9a4ec188a8d32caf5526ec8ee335b4ee5b
1 parent 29336b9 commit 3f1b6d8

4 files changed

Lines changed: 850 additions & 0 deletions

File tree

‎cinderx/benchmarks/torchrec_pt2/README.md‎

Lines changed: 62 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,62 @@
1+
# CinderX + TorchRec pt2_compile_benchmark (CPU)
2+
3+
Benchmarks CinderX JIT impact on the TorchRec `pt2_compile_benchmark`
4+
(eager forward+backward of DistributedModelParallel), adapted to run on CPU.
5+
Uses `FakeStore` and a fake process group — no GPU or multi-node setup needed.
6+
7+
## Running the Benchmark
8+
9+
### CinderX JIT benchmark
10+
11+
```bash
12+
# Quick test (5 features, default)
13+
buck run //cinderx/benchmarks/torchrec_pt2:run_with_cinderx -- --num-features 5 --repeat 1 --iters 10
14+
15+
# Full benchmark (200 features, 50 iters x 3 runs)
16+
buck run //cinderx/benchmarks/torchrec_pt2:run_with_cinderx -- --num-features 200 --repeat 3 --iters 50
17+
```
18+
19+
### Without JIT (for comparison)
20+
21+
Set `PYTHONJITDISABLE=1` to run the same binary with the JIT forcefully disabled:
22+
23+
```bash
24+
PYTHONJITDISABLE=1 buck run //cinderx/benchmarks/torchrec_pt2:run_with_cinderx -- \
25+
--num-features 200 --repeat 3 --iters 50
26+
```
27+
28+
### Baseline (original compile benchmark, no CinderX)
29+
30+
```bash
31+
buck run //cinderx/benchmarks/torchrec_pt2:baseline -- --num-features 200
32+
```
33+
34+
### Specific Python version
35+
36+
The `cinderx_benchmark_binary` macro generates versioned targets:
37+
38+
```bash
39+
buck run //cinderx/benchmarks/torchrec_pt2:run_with_cinderx-314 -- --num-features 200
40+
buck run //cinderx/benchmarks/torchrec_pt2:run_with_cinderx-312 -- --num-features 200
41+
```
42+
43+
## Files
44+
45+
| File | Description |
46+
|---|---|
47+
| `run_with_cinderx.py` | CinderX JIT benchmark runner |
48+
| `test_pt2_multiprocess.py` | Adapted benchmark (CPU, baseline model/sharders) |
49+
| `BUCK` | Build targets |
50+
| `README.md` | This file |
51+
52+
## CPU Adaptations
53+
54+
The original benchmark (`fbcode/torchrec/distributed/tests/test_pt2_multiprocess.py`)
55+
runs on GPU. This version is adapted to run on CPU so it can run anywhere without
56+
GPU or multi-node setup:
57+
58+
1. `torch.device("cuda")` -> `torch.device("cpu")`
59+
2. `compute_device="cuda"` -> `compute_device="cpu"` in Topology
60+
3. `fullgraph=True` -> `fullgraph=False` (fbgemm CPU ops don't support full graph)
61+
4. `convert_to_vb=True` -> `convert_to_vb=False` (VBE path uses unsupported CPU ops)
62+
5. Relaxed eager-vs-compile assertion (CPU fake PG produces NaN in compile path)

‎cinderx/benchmarks/torchrec_pt2/run_with_cinderx.py‎

Lines changed: 298 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,298 @@
1+
#!/usr/bin/env python3
2+
# Copyright (c) Meta Platforms, Inc. and affiliates.
3+
# All rights reserved.
4+
#
5+
# This source code is licensed under the BSD-style license found in the
6+
# LICENSE file in the root directory of this source tree.
7+
8+
"""Run the TorchRec DMP eager forward+backward with CinderX JIT enabled.
9+
10+
Measures CinderX JIT auto-compilation impact on the Python-level orchestration
11+
code in torchrec's distributed embedding sharding (DistributedModelParallel).
12+
13+
Usage (via buck):
14+
buck run //cinderx/benchmarks/torchrec_pt2:run_with_cinderx
15+
16+
# With options:
17+
buck run //cinderx/benchmarks/torchrec_pt2:run_with_cinderx -- \\
18+
--num-features 200 --repeat 3 --iters 50
19+
20+
# Without JIT (for comparison):
21+
PYTHONJITDISABLE=1 buck run //cinderx/benchmarks/torchrec_pt2:run_with_cinderx -- \\
22+
--num-features 200 --repeat 3 --iters 50
23+
24+
# Specific Python version:
25+
buck run //cinderx/benchmarks/torchrec_pt2:run_with_cinderx-314 -- --num-features 200
26+
"""
27+
28+
import sys
29+
import time
30+
31+
_t0 = time.perf_counter()
32+
from cinderx.compiler.strict import loader as static_python_loader
33+
34+
static_python_loader.install()
35+
_static_loader_time = time.perf_counter() - _t0
36+
37+
import _static
38+
from cinderx.benchmarks.torchrec_pt2.static_helper import add_one
39+
40+
assert _static.is_static_callable(add_one), (
41+
"static_helper.add_one was not statically compiled! "
42+
f"co_flags=0x{add_one.__code__.co_flags:x}"
43+
)
44+
print(f"Static compilation check passed: add_one is statically compiled")
45+
print(f"add_one(41) = {add_one(41)}")
46+
47+
import cinderx.jit
48+
import click
49+
import torch
50+
import torch._dynamo
51+
import torchrec
52+
import torchrec.pt2.checks
53+
from cinderx.benchmarks.torchrec_pt2.test_pt2_multiprocess import (
54+
_gen_model,
55+
_ModelType,
56+
_TestConfig,
57+
EBCSharderFixedShardingType,
58+
ECSharderFixedShardingType,
59+
TestModelInfo,
60+
)
61+
from torch import distributed as dist
62+
from torch._dynamo.testing import reduce_to_scalar_loss
63+
from torch.distributed import ProcessGroup
64+
from torch.testing._internal.distributed.fake_pg import FakeStore
65+
from torchrec.distributed.embedding import EmbeddingCollectionSharder
66+
from torchrec.distributed.model_parallel import DistributedModelParallel
67+
from torchrec.distributed.planner import EmbeddingShardingPlanner, Topology
68+
from torchrec.distributed.planner.enumerators import EmbeddingEnumerator
69+
from torchrec.distributed.planner.shard_estimators import (
70+
EmbeddingPerfEstimator,
71+
EmbeddingStorageEstimator,
72+
)
73+
from torchrec.distributed.planner.types import ShardingPlan
74+
from torchrec.distributed.sharding_plan import EmbeddingBagCollectionSharder
75+
from torchrec.distributed.test_utils.test_model import ModelInput
76+
from torchrec.distributed.types import ShardingEnv, ShardingType
77+
from torchrec.modules.embedding_modules import (
78+
EmbeddingBagCollection,
79+
EmbeddingBagConfig,
80+
)
81+
from torchrec.pt2.utils import kjt_for_pt2_tracing
82+
from torchrec.sparse.jagged_tensor import KeyedJaggedTensor, KeyedTensor
83+
84+
85+
def setup_benchmark(
86+
rank: int = 0,
87+
world_size: int = 2,
88+
num_features: int = 5,
89+
batch_size: int = 10,
90+
num_embeddings: int = 256,
91+
):
92+
"""Set up the benchmark (model, DMP, inputs). Returns everything needed."""
93+
sharding_type = ShardingType.TABLE_WISE.value
94+
emb_dim = 12
95+
num_float_features: int = 8
96+
num_weighted_features: int = 1
97+
98+
device: torch.Device = torch.device("cpu")
99+
store = FakeStore()
100+
dist.init_process_group(
101+
backend="fake", rank=rank, world_size=world_size, store=store
102+
)
103+
pg: ProcessGroup = dist.distributed_c10d._get_default_group()
104+
105+
topology: Topology = Topology(world_size=world_size, compute_device="cpu")
106+
mi = TestModelInfo(
107+
# pyrefly: ignore [bad-argument-type]
108+
dense_device=device,
109+
# pyrefly: ignore [bad-argument-type]
110+
sparse_device=device,
111+
num_features=num_features,
112+
num_float_features=num_float_features,
113+
num_weighted_features=num_weighted_features,
114+
topology=topology,
115+
)
116+
117+
mi.planner = EmbeddingShardingPlanner(
118+
topology=topology,
119+
batch_size=batch_size,
120+
enumerator=EmbeddingEnumerator(
121+
topology=topology,
122+
batch_size=batch_size,
123+
estimator=[
124+
EmbeddingPerfEstimator(topology=topology),
125+
EmbeddingStorageEstimator(topology=topology),
126+
],
127+
),
128+
)
129+
130+
mi.tables = [
131+
EmbeddingBagConfig(
132+
num_embeddings=num_embeddings,
133+
embedding_dim=emb_dim,
134+
name="table_" + str(i),
135+
feature_names=["feature_" + str(i)],
136+
)
137+
for i in range(mi.num_features)
138+
]
139+
140+
mi.weighted_tables = [
141+
EmbeddingBagConfig(
142+
num_embeddings=num_embeddings,
143+
embedding_dim=emb_dim,
144+
name="weighted_table_" + str(i),
145+
feature_names=["weighted_feature_" + str(i)],
146+
)
147+
for i in range(mi.num_weighted_features)
148+
]
149+
150+
mi.model = _gen_model(_ModelType.EBC, mi)
151+
mi.model.training = True
152+
153+
planner = EmbeddingShardingPlanner(
154+
topology=Topology(world_size, device.type),
155+
constraints=None,
156+
)
157+
158+
sharders = [
159+
EBCSharderFixedShardingType(sharding_type),
160+
ECSharderFixedShardingType(sharding_type),
161+
]
162+
163+
plan: ShardingPlan = planner.plan(mi.model, sharders)
164+
165+
dmp = DistributedModelParallel(
166+
mi.model,
167+
env=ShardingEnv(world_size, rank, pg),
168+
plan=plan,
169+
# pyrefly: ignore [bad-argument-type]
170+
sharders=sharders,
171+
# pyrefly: ignore [bad-argument-type]
172+
device=device,
173+
init_data_parallel=False,
174+
)
175+
176+
_, local_model_inputs = ModelInput.generate(
177+
batch_size=batch_size,
178+
world_size=world_size,
179+
num_float_features=num_float_features,
180+
tables=mi.tables,
181+
weighted_tables=mi.weighted_tables,
182+
variable_batch_size=False,
183+
)
184+
185+
# pyrefly: ignore [bad-argument-type]
186+
local_model_input = local_model_inputs[rank].to(device)
187+
kjt = local_model_input.idlist_features
188+
ff = local_model_input.float_features
189+
ff.requires_grad = True
190+
# pyrefly: ignore [bad-argument-type]
191+
kjt_ft = kjt_for_pt2_tracing(kjt, convert_to_vb=False)
192+
193+
if hasattr(torchrec.distributed, "comm_ops") and hasattr(
194+
# pyrefly: ignore [implicit-import]
195+
torchrec.distributed.comm_ops,
196+
"set_use_sync_collectives",
197+
):
198+
# pyrefly: ignore [implicit-import]
199+
torchrec.distributed.comm_ops.set_use_sync_collectives(True)
200+
torchrec.pt2.checks.set_use_torchdynamo_compiling_path(True)
201+
202+
dmp.train(True)
203+
204+
return {
205+
"dmp": dmp,
206+
"kjt_ft": kjt_ft,
207+
"ff": ff,
208+
}
209+
210+
211+
def run_eager_iters(ctx: dict, num_iters: int) -> float:
212+
"""Run timed eager forward+backward iterations. Returns elapsed seconds."""
213+
dmp = ctx["dmp"]
214+
kjt_ft = ctx["kjt_ft"]
215+
ff = ctx["ff"]
216+
217+
t_start = time.perf_counter()
218+
for _ in range(num_iters):
219+
out = dmp(kjt_ft, ff)
220+
reduce_to_scalar_loss(out).backward()
221+
t_end = time.perf_counter()
222+
return t_end - t_start
223+
224+
225+
@click.command()
226+
@click.option("--repeat", type=int, default=3, help="Number of timed runs")
227+
@click.option(
228+
"--iters", type=int, default=50, help="Forward+backward iterations per run"
229+
)
230+
@click.option("--warmup", type=int, default=5, help="Warmup iterations (not timed)")
231+
@click.option("--rank", type=int, default=0)
232+
@click.option("--world-size", type=int, default=2)
233+
@click.option("--num-features", type=int, default=200)
234+
@click.option("--batch-size", type=int, default=10)
235+
def main(
236+
rank: int,
237+
world_size: int,
238+
num_features: int,
239+
batch_size: int,
240+
repeat: int,
241+
iters: int,
242+
warmup: int,
243+
):
244+
print(f"Python {sys.version}")
245+
print("CinderX JIT eager forward+backward benchmark")
246+
print(f"Static loader import+install: {_static_loader_time * 1000:.1f}ms")
247+
print(
248+
f"num_features={num_features}, batch_size={batch_size}, "
249+
f"warmup={warmup}, iters={iters}, repeat={repeat}"
250+
)
251+
print()
252+
253+
# Setup model and DMP
254+
print("Setting up model and DMP...")
255+
ctx = setup_benchmark(
256+
rank=rank,
257+
world_size=world_size,
258+
num_features=num_features,
259+
batch_size=batch_size,
260+
)
261+
262+
# Enable JIT auto-compilation
263+
print("Enabling CinderX JIT...")
264+
cinderx.jit.enable()
265+
cinderx.jit.auto()
266+
267+
# Warmup (lets JIT compile hot functions)
268+
print(f"Warmup ({warmup} iters)...")
269+
run_eager_iters(ctx, warmup)
270+
271+
# Timed runs
272+
times = []
273+
for i in range(repeat):
274+
t = run_eager_iters(ctx, iters)
275+
times.append(t)
276+
print(
277+
f" Run {i + 1}/{repeat}: {iters} iters in {t:.3f}s "
278+
f"({t / iters * 1000:.1f}ms/iter)"
279+
)
280+
281+
cinderx.jit.disable()
282+
dist.destroy_process_group()
283+
284+
avg = sum(times) / len(times)
285+
avg_per_iter = avg / iters
286+
print()
287+
print("=" * 60)
288+
print("CinderX JIT results (eager forward+backward):")
289+
print(f" num_features={num_features}, batch_size={batch_size}")
290+
print(f" warmup={warmup}, iters_per_run={iters}, runs={repeat}")
291+
print(f" Run times: {[f'{t:.3f}s' for t in times]}")
292+
print(f" Avg per run: {avg:.3f}s")
293+
print(f" Avg per iter: {avg_per_iter * 1000:.1f}ms")
294+
print("=" * 60)
295+
296+
297+
if __name__ == "__main__":
298+
main()

‎cinderx/benchmarks/torchrec_pt2/static_helper.py‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,7 @@
1+
# (c) Meta Platforms, Inc. and affiliates. Confidential and proprietary.
2+
3+
import __static__
4+
5+
6+
def add_one(x: int) -> int:
7+
return x + 1

0 commit comments

Comments
 (0)