Skip to content

Commit 2e2b27c

Browse files
pagezyhfmollyheamazon
pagezyhf
and
mollyheamazon
authored
[hf] HF PT Training DLCs (#5301)
* image * add py312 * fix * test fix * typo --------- Co-authored-by: Molly He <[email protected]>
1 parent 1d3f789 commit 2e2b27c

File tree

3 files changed

+104
-1
lines changed

3 files changed

+104
-1
lines changed

‎src/sagemaker/fw_utils.py‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -156,6 +156,9 @@
156156
"2.3.1",
157157
"2.4.1",
158158
"2.5.1",
159+
"2.6.0",
160+
"2.7.1",
161+
"2.8.0",
159162
]
160163

161164
TRAINIUM_SUPPORTED_DISTRIBUTION_STRATEGIES = ["torch_distributed"]

‎src/sagemaker/image_uri_config/huggingface.json‎

Lines changed: 97 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,9 @@
1616
"4.36": "4.36.0",
1717
"4.46": "4.46.1",
1818
"4.48": "4.48.0",
19-
"4.49": "4.49.0"
19+
"4.49": "4.49.0",
20+
"4.55": "4.55.0",
21+
"4.56": "4.56.2"
2022
},
2123
"versions": {
2224
"4.4.2": {
@@ -1162,6 +1164,100 @@
11621164
"gpu": "cu124-ubuntu22.04"
11631165
}
11641166
}
1167+
},
1168+
"4.55.0": {
1169+
"version_aliases": {
1170+
"pytorch2.7": "pytorch2.7.1"
1171+
},
1172+
"pytorch2.7.1": {
1173+
"py_versions": [
1174+
"py312"
1175+
],
1176+
"registries": {
1177+
"af-south-1": "626614931356",
1178+
"il-central-1": "780543022126",
1179+
"ap-east-1": "871362719292",
1180+
"ap-northeast-1": "763104351884",
1181+
"ap-northeast-2": "763104351884",
1182+
"ap-northeast-3": "364406365360",
1183+
"ap-south-1": "763104351884",
1184+
"ap-southeast-1": "763104351884",
1185+
"ap-southeast-2": "763104351884",
1186+
"ap-southeast-3": "907027046896",
1187+
"ca-central-1": "763104351884",
1188+
"cn-north-1": "727897471807",
1189+
"cn-northwest-1": "727897471807",
1190+
"eu-central-1": "763104351884",
1191+
"eu-north-1": "763104351884",
1192+
"eu-west-1": "763104351884",
1193+
"eu-west-2": "763104351884",
1194+
"eu-west-3": "763104351884",
1195+
"eu-south-1": "692866216735",
1196+
"me-south-1": "217643126080",
1197+
"me-central-1": "914824155844",
1198+
"sa-east-1": "763104351884",
1199+
"us-east-1": "763104351884",
1200+
"us-east-2": "763104351884",
1201+
"us-gov-east-1": "446045086412",
1202+
"us-gov-west-1": "442386744353",
1203+
"us-iso-east-1": "886529160074",
1204+
"us-isob-east-1": "094389454867",
1205+
"us-west-1": "763104351884",
1206+
"us-west-2": "763104351884",
1207+
"ca-west-1": "204538143572"
1208+
},
1209+
"repository": "huggingface-pytorch-training",
1210+
"container_version": {
1211+
"gpu": "cu128-ubuntu22.04"
1212+
}
1213+
}
1214+
},
1215+
"4.56.2": {
1216+
"version_aliases": {
1217+
"pytorch2.8": "pytorch2.8.0"
1218+
},
1219+
"pytorch2.8.0": {
1220+
"py_versions": [
1221+
"py312"
1222+
],
1223+
"registries": {
1224+
"af-south-1": "626614931356",
1225+
"il-central-1": "780543022126",
1226+
"ap-east-1": "871362719292",
1227+
"ap-northeast-1": "763104351884",
1228+
"ap-northeast-2": "763104351884",
1229+
"ap-northeast-3": "364406365360",
1230+
"ap-south-1": "763104351884",
1231+
"ap-southeast-1": "763104351884",
1232+
"ap-southeast-2": "763104351884",
1233+
"ap-southeast-3": "907027046896",
1234+
"ca-central-1": "763104351884",
1235+
"cn-north-1": "727897471807",
1236+
"cn-northwest-1": "727897471807",
1237+
"eu-central-1": "763104351884",
1238+
"eu-north-1": "763104351884",
1239+
"eu-west-1": "763104351884",
1240+
"eu-west-2": "763104351884",
1241+
"eu-west-3": "763104351884",
1242+
"eu-south-1": "692866216735",
1243+
"me-south-1": "217643126080",
1244+
"me-central-1": "914824155844",
1245+
"sa-east-1": "763104351884",
1246+
"us-east-1": "763104351884",
1247+
"us-east-2": "763104351884",
1248+
"us-gov-east-1": "446045086412",
1249+
"us-gov-west-1": "442386744353",
1250+
"us-iso-east-1": "886529160074",
1251+
"us-isob-east-1": "094389454867",
1252+
"us-west-1": "763104351884",
1253+
"us-west-2": "763104351884",
1254+
"ca-west-1": "204538143572"
1255+
},
1256+
"repository": "huggingface-pytorch-training",
1257+
"container_version": {
1258+
"gpu": "cu129-ubuntu22.04"
1259+
}
1260+
}
11651261
}
11661262
}
11671263
},

‎tests/conftest.py‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -297,6 +297,8 @@ def huggingface_pytorch_training_version(huggingface_training_version):
297297

298298
@pytest.fixture(scope="module")
299299
def huggingface_pytorch_training_py_version(huggingface_pytorch_training_version):
300+
if Version(huggingface_pytorch_training_version) >= Version("2.6"):
301+
return "py312"
300302
if Version(huggingface_pytorch_training_version) >= Version("2.3"):
301303
return "py311"
302304
if Version(huggingface_pytorch_training_version) >= Version("2.0"):
@@ -361,6 +363,8 @@ def huggingface_training_compiler_pytorch_py_version(
361363
def huggingface_pytorch_latest_training_py_version(
362364
huggingface_training_pytorch_latest_version,
363365
):
366+
if Version(huggingface_training_pytorch_latest_version) >= Version("2.6"):
367+
return "py312"
364368
if Version(huggingface_training_pytorch_latest_version) >= Version("2.3"):
365369
return "py311"
366370
if Version(huggingface_training_pytorch_latest_version) >= Version("2.0"):

0 commit comments

Comments
 (0)