Skip to content

Commit 95f500a

Browse files
qiurcfacebook-github-bot
authored andcommitted
Add option to support garbage collection after torch compilation (#2559)
Summary: X-link: pytorch/pytorch#142821 This diff is an extension of ezyang's PR https://fburl.com/6uvvzb4f. In ezyang's PR above, it adds gc after torch compilation finished. The gc operation is guarded by jk: pytorch/compiler:enable_run_gc_after_compile The gc op time cost will be logged into dynamo_compile scuba table. This diff extends the PR to: - Use garbage collection on Generation 1 instead of generation 2 (default), which greatly reduced the gc latency overhead from 160 sec per rank to 10 sec per rank. - Additionally introduce an environment variance which has the higher priority than the JK to control whether we do gc or not after the torch compilation. (default value set to gc enabled). This environment variance will be used for AB testing of training jobs to compare the pt2 compilation time and memory cost. Reviewed By: ezyang, yuxihu Differential Revision: D67062158
1 parent 6f191e9 commit 95f500a

File tree

1 file changed

+1
-0
lines changed
  • userbenchmark/dynamo/dynamobench/_dynamo

1 file changed

+1
-0
lines changed

userbenchmark/dynamo/dynamobench/_dynamo/utils.py

+1
Original file line numberDiff line numberDiff line change
@@ -882,6 +882,7 @@ class CompilationMetrics:
882882
triton_version: Optional[str] = None
883883
feature_usage: Optional[dict[str, bool]] = None
884884
compile_time_autotune_time_us: Optional[int] = None
885+
gc_time_us: Optional[int] = None
885886

886887

887888
DEFAULT_COMPILATION_METRICS_LIMIT = 64

0 commit comments

Comments
 (0)