frozenleaves/pytorch
1
0
Fork 0
mirror of https://github.com/pytorch/pytorch.git synced 2025-10-20 21:14:14 +08:00
Code Issues Packages Projects Releases Wiki Activity

Combine sticky pgo key with job id (#154863)

Browse Source 
Pull Request resolved: https://github.com/pytorch/pytorch/pull/154863
Approved by: https://github.com/Mingming-Ding
This commit is contained in:
bobrenjc93
2025-06-02 13:57:40 -07:00
committed by PyTorch MergeBot
parent a4da1d4a47
commit ea5b9eca74
2 changed files with 4 additions and 13 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
torch/_dynamo/pgo.py
View File

@ -535,9 +535,6 @@ def get_cache_key() -> Optional[str]:
)
return f"{r}:{rank}:{tag}"
if r := torch.compiler.config.sticky_pgo_key:
return f"sticky:{r}:{rank}:{tag}"
if (name_version := torch._utils_internal.get_mast_job_name_version()) is not None:
mast_job_name, mast_job_version = name_version
return f"mast:{mast_job_name}:{mast_job_version}:{rank}:{tag}"

14
torch/compiler/config.py
View File

@ -29,7 +29,10 @@ __all__ = [
# FB-internal note: you do NOT have to specify this explicitly specify this if
# you run on MAST, we will automatically default this to
# mast:MAST_JOB_NAME:MAST_JOB_VERSION.
job_id: Optional[str] = Config(env_name_default="TORCH_COMPILE_JOB_ID", default=None)
job_id: Optional[str] = Config(
env_name_default=["TORCH_COMPILE_JOB_ID", "TORCH_COMPILE_STICKY_PGO_KEY"],
default=None,
)
"""
Semantically, this should be an identifier that uniquely identifies, e.g., a
training job. You might have multiple attempts of the same job, e.g., if it was
@ -74,15 +77,6 @@ This whitelist is dominant over all other flags dynamic=False, force_nn_module_p
and force_parameter_static_shapes.
"""
sticky_pgo_key: str = Config(
env_name_default="TORCH_COMPILE_STICKY_PGO_KEY", default=""
)
"""
If you want to share PGO profiles across different jobs (and not just attempts), you can set
this to a string that identifies the shared profile. This is useful if you want to share PGO profiles
for models that are not identical, but are similar enough to share PGO profiles.
"""
unbacked_sources: str = Config(
env_name_default="TORCH_COMPILE_UNBACKED_SOURCES", default=""
)