From f48e8435d4616307b382cd638d979194a01cf72a Mon Sep 17 00:00:00 2001 From: Arjun Krishnakumar Date: Wed, 16 Sep 2026 15:44:29 +0200 Subject: [PATCH 1/5] chore: pin datasets to 5.0.0 in the base dependencies --- pyproject.toml | 2 +- uv.lock | 90 ++++++++++++++++++++++++++++++++++---------------- 2 files changed, 63 insertions(+), 29 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index b2e0d43..cbffae4 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -13,6 +13,7 @@ dependencies = [ "wandb", "huggingface-hub", "psutil", + "datasets==5.0.0", ] [project.optional-dependencies] @@ -21,7 +22,6 @@ trl = [ "trl", "deepspeed", "transformers", - "datasets>4.5.0", "accelerate", "kernels", "flash_attn", diff --git a/uv.lock b/uv.lock index 22ad5b9..819586e 100644 --- a/uv.lock +++ b/uv.lock @@ -162,7 +162,7 @@ wheels = [ [[package]] name = "datasets" -version = "4.6.1" +version = "5.0.0" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "dill", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, @@ -180,9 +180,9 @@ dependencies = [ { name = "tqdm", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, { name = "xxhash", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/d7/94/eb81c6fe32e9b6ef92223141b5a553aeff2e9456968424a8533cbe88f476/datasets-4.6.1.tar.gz", hash = "sha256:140ce500bc41939ff6ce995702d66b1f4b2ee7f117bb9b07512fab6804d4070a", size = 593865, upload-time = "2026-02-27T23:26:49.482Z" } +sdist = { url = "https://files.pythonhosted.org/packages/d9/85/ce4f780c32f7e36d71257f1c27e8ba898ebe379cb54f211f5f2013f2c219/datasets-5.0.0.tar.gz", hash = "sha256:83dbbbdb07a33b82192b8c419deb18739b138ee2ce1a322d55ce6b100954ec1a", size = 631708, upload-time = "2026-06-05T13:18:26.124Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/37/f0/99fe6eb530c7ee9ee1faee48059eb8a6437f80c893a496b98a78864e0fc6/datasets-4.6.1-py3-none-any.whl", hash = "sha256:f53228e6dadc9f837037b1bf3051d7d8c054abbb3eb29f1f022926e08090e0da", size = 520667, upload-time = "2026-02-27T23:26:46.855Z" }, + { url = "https://files.pythonhosted.org/packages/05/66/73034ad30b59f13439b75e620989dacba4c047256e358ba7c2e9ec98ea22/datasets-5.0.0-py3-none-any.whl", hash = "sha256:7dd34927a0fd7046e98aad5cb9430e699c373238a15befa7b9bf22b991a7fee6", size = 555084, upload-time = "2026-06-05T13:18:24.435Z" }, ] [[package]] @@ -234,10 +234,10 @@ wheels = [ [[package]] name = "filelock" version = "3.25.0" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/77/18/a1fd2231c679dcb9726204645721b12498aeac28e1ad0601038f94b42556/filelock-3.25.0.tar.gz", hash = "sha256:8f00faf3abf9dc730a1ffe9c354ae5c04e079ab7d3a683b7c32da5dd05f26af3", size = 40158, upload-time = "2026-03-01T15:08:45.916Z" } +source = { registry = "https://download.pytorch.org/whl/cu126" } +sdist = { url = "https://files.pythonhosted.org/packages/77/18/a1fd2231c679dcb9726204645721b12498aeac28e1ad0601038f94b42556/filelock-3.25.0.tar.gz", hash = "sha256:8f00faf3abf9dc730a1ffe9c354ae5c04e079ab7d3a683b7c32da5dd05f26af3" } wheels = [ - { url = "https://files.pythonhosted.org/packages/f9/0b/de6f54d4a8bedfe8645c41497f3c18d749f0bd3218170c667bf4b81d0cdd/filelock-3.25.0-py3-none-any.whl", hash = "sha256:5ccf8069f7948f494968fc0713c10e5c182a9c9d9eef3a636307a20c2490f047", size = 26427, upload-time = "2026-03-01T15:08:44.593Z" }, + { url = "https://files.pythonhosted.org/packages/f9/0b/de6f54d4a8bedfe8645c41497f3c18d749f0bd3218170c667bf4b81d0cdd/filelock-3.25.0-py3-none-any.whl", hash = "sha256:5ccf8069f7948f494968fc0713c10e5c182a9c9d9eef3a636307a20c2490f047" }, ] [[package]] @@ -272,10 +272,10 @@ wheels = [ [[package]] name = "fsspec" version = "2026.2.0" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/51/7c/f60c259dcbf4f0c47cc4ddb8f7720d2dcdc8888c8e5ad84c73ea4531cc5b/fsspec-2026.2.0.tar.gz", hash = "sha256:6544e34b16869f5aacd5b90bdf1a71acb37792ea3ddf6125ee69a22a53fb8bff", size = 313441, upload-time = "2026-02-05T21:50:53.743Z" } +source = { registry = "https://download.pytorch.org/whl/cu126" } +sdist = { url = "https://files.pythonhosted.org/packages/51/7c/f60c259dcbf4f0c47cc4ddb8f7720d2dcdc8888c8e5ad84c73ea4531cc5b/fsspec-2026.2.0.tar.gz", hash = "sha256:6544e34b16869f5aacd5b90bdf1a71acb37792ea3ddf6125ee69a22a53fb8bff" } wheels = [ - { url = "https://files.pythonhosted.org/packages/e6/ab/fb21f4c939bb440104cc2b396d3be1d9b7a9fd3c6c2a53d98c45b3d7c954/fsspec-2026.2.0-py3-none-any.whl", hash = "sha256:98de475b5cb3bd66bedd5c4679e87b4fdfe1a3bf4d707b151b3c07e58c9a2437", size = 202505, upload-time = "2026-02-05T21:50:51.819Z" }, + { url = "https://files.pythonhosted.org/packages/e6/ab/fb21f4c939bb440104cc2b396d3be1d9b7a9fd3c6c2a53d98c45b3d7c954/fsspec-2026.2.0-py3-none-any.whl", hash = "sha256:98de475b5cb3bd66bedd5c4679e87b4fdfe1a3bf4d707b151b3c07e58c9a2437" }, ] [package.optional-dependencies] @@ -414,6 +414,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/0e/61/66938bbb5fc52dbdf84594873d5b51fb1f7c7794e9c0f5bd885f30bc507b/idna-3.11-py3-none-any.whl", hash = "sha256:771a87f49d9defaf64091e6e6fe9c18d4833f140bd19464795bc32d966ca37ea", size = 71008, upload-time = "2025-10-12T14:55:18.883Z" }, ] +[[package]] +name = "iniconfig" +version = "2.3.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/72/34/14ca021ce8e5dfedc35312d08ba8bf51fdd999c576889fc2c24cb97f4f10/iniconfig-2.3.0.tar.gz", hash = "sha256:c76315c77db068650d49c5b56314774a7804df16fee4402c1f19d6d15d8c4730", size = 20503, upload-time = "2025-10-18T21:55:43.219Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/cb/b1/3846dd7f199d53cb17f49cba7e651e9ce294d8497c8c150530ed11865bb8/iniconfig-2.3.0-py3-none-any.whl", hash = "sha256:f631c04d2c48c52b84d0d0549c99ff3859c98df65b3101406327ecc7d53fbf12", size = 7484, upload-time = "2025-10-18T21:55:41.639Z" }, +] + [[package]] name = "jinja2" version = "3.1.6" @@ -422,7 +431,7 @@ dependencies = [ { name = "markupsafe", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, ] wheels = [ - { url = "https://download.pytorch.org/whl/jinja2-3.1.6-py3-none-any.whl" }, + { url = "https://download.pytorch.org/whl/jinja2-3.1.6-py3-none-any.whl", upload-time = "2025-10-14T18:38:59Z" }, ] [[package]] @@ -476,11 +485,9 @@ wheels = [ [[package]] name = "markupsafe" version = "3.0.3" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/7e/99/7690b6d4034fffd95959cbe0c02de8deb3098cc577c67bb6a24fe5d7caa7/markupsafe-3.0.3.tar.gz", hash = "sha256:722695808f4b6457b320fdc131280796bdceb04ab50fe1795cd540799ebe1698", size = 80313, upload-time = "2025-09-27T18:37:40.426Z" } +source = { registry = "https://download.pytorch.org/whl/cu126" } wheels = [ - { url = "https://files.pythonhosted.org/packages/3c/2e/8d0c2ab90a8c1d9a24f0399058ab8519a3279d1bd4289511d74e909f060e/markupsafe-3.0.3-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:d6dd0be5b5b189d31db7cda48b91d7e0a9795f31430b7f271219ab30f1d3ac9d", size = 22947, upload-time = "2025-09-27T18:36:33.86Z" }, - { url = "https://files.pythonhosted.org/packages/89/e0/4486f11e51bbba8b0c041098859e869e304d1c261e59244baa3d295d47b7/markupsafe-3.0.3-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:77f0643abe7495da77fb436f50f8dab76dbc6e5fd25d39589a0f1fe6548bfa2b", size = 23015, upload-time = "2025-09-27T18:36:37.868Z" }, + { url = "https://download.pytorch.org/whl/markupsafe-3.0.3-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", upload-time = "2026-03-27T13:54:35Z" }, ] [[package]] @@ -569,11 +576,11 @@ wheels = [ [[package]] name = "numpy" version = "2.4.2" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/57/fd/0005efbd0af48e55eb3c7208af93f2862d4b1a56cd78e84309a2d959208d/numpy-2.4.2.tar.gz", hash = "sha256:659a6107e31a83c4e33f763942275fd278b21d095094044eb35569e86a21ddae", size = 20723651, upload-time = "2026-01-31T23:13:10.135Z" } +source = { registry = "https://download.pytorch.org/whl/cu126" } +sdist = { url = "https://files.pythonhosted.org/packages/57/fd/0005efbd0af48e55eb3c7208af93f2862d4b1a56cd78e84309a2d959208d/numpy-2.4.2.tar.gz", hash = "sha256:659a6107e31a83c4e33f763942275fd278b21d095094044eb35569e86a21ddae" } wheels = [ - { url = "https://files.pythonhosted.org/packages/f5/c6/a18e59f3f0b8071cc85cbc8d80cd02d68aa9710170b2553a117203d46936/numpy-2.4.2-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:9e35d3e0144137d9fdae62912e869136164534d64a169f86438bc9561b6ad49f", size = 16619782, upload-time = "2026-01-31T23:11:03.669Z" }, - { url = "https://files.pythonhosted.org/packages/61/c4/c4066322256ec740acc1c8923a10047818691d2f8aec254798f3dd90f5f2/numpy-2.4.2-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:8906e71fd8afcb76580404e2a950caef2685df3d2a57fe82a86ac8d33cc007ba", size = 18345324, upload-time = "2026-01-31T23:11:08.248Z" }, + { url = "https://files.pythonhosted.org/packages/f5/c6/a18e59f3f0b8071cc85cbc8d80cd02d68aa9710170b2553a117203d46936/numpy-2.4.2-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:9e35d3e0144137d9fdae62912e869136164534d64a169f86438bc9561b6ad49f" }, + { url = "https://files.pythonhosted.org/packages/61/c4/c4066322256ec740acc1c8923a10047818691d2f8aec254798f3dd90f5f2/numpy-2.4.2-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:8906e71fd8afcb76580404e2a950caef2685df3d2a57fe82a86ac8d33cc007ba" }, ] [[package]] @@ -756,12 +763,12 @@ wheels = [ [[package]] name = "pillow" version = "12.1.1" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/1f/42/5c74462b4fd957fcd7b13b04fb3205ff8349236ea74c7c375766d6c82288/pillow-12.1.1.tar.gz", hash = "sha256:9ad8fa5937ab05218e2b6a4cff30295ad35afd2f83ac592e68c0d871bb0fdbc4", size = 46980264, upload-time = "2026-02-11T04:23:07.146Z" } +source = { registry = "https://download.pytorch.org/whl/cu126" } +sdist = { url = "https://files.pythonhosted.org/packages/1f/42/5c74462b4fd957fcd7b13b04fb3205ff8349236ea74c7c375766d6c82288/pillow-12.1.1.tar.gz", hash = "sha256:9ad8fa5937ab05218e2b6a4cff30295ad35afd2f83ac592e68c0d871bb0fdbc4" } wheels = [ - { url = "https://files.pythonhosted.org/packages/07/26/246ab11455b2549b9233dbd44d358d033a2f780fa9007b61a913c5b2d24e/pillow-12.1.1-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:aee2810642b2898bb187ced9b349e95d2a7272930796e022efaf12e99dccd293", size = 8045012, upload-time = "2026-02-11T04:20:52.882Z" }, - { url = "https://files.pythonhosted.org/packages/ff/79/6df7b2ee763d619cda2fb4fea498e5f79d984dae304d45a8999b80d6cf5c/pillow-12.1.1-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:7aac39bcf8d4770d089588a2e1dd111cbaa42df5a94be3114222057d68336bd0", size = 7041540, upload-time = "2026-02-11T04:20:55.97Z" }, - { url = "https://files.pythonhosted.org/packages/03/03/31216ec124bb5c3dacd74ce8efff4cc7f52643653bad4825f8f08c697743/pillow-12.1.1-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:339ffdcb7cbeaa08221cd401d517d4b1fe7a9ed5d400e4a8039719238620ca35", size = 7166745, upload-time = "2026-02-11T04:20:59.196Z" }, + { url = "https://files.pythonhosted.org/packages/07/26/246ab11455b2549b9233dbd44d358d033a2f780fa9007b61a913c5b2d24e/pillow-12.1.1-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:aee2810642b2898bb187ced9b349e95d2a7272930796e022efaf12e99dccd293" }, + { url = "https://files.pythonhosted.org/packages/ff/79/6df7b2ee763d619cda2fb4fea498e5f79d984dae304d45a8999b80d6cf5c/pillow-12.1.1-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:7aac39bcf8d4770d089588a2e1dd111cbaa42df5a94be3114222057d68336bd0" }, + { url = "https://files.pythonhosted.org/packages/03/03/31216ec124bb5c3dacd74ce8efff4cc7f52643653bad4825f8f08c697743/pillow-12.1.1-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:339ffdcb7cbeaa08221cd401d517d4b1fe7a9ed5d400e4a8039719238620ca35" }, ] [[package]] @@ -773,11 +780,21 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/48/31/05e764397056194206169869b50cf2fee4dbbbc71b344705b9c0d878d4d8/platformdirs-4.9.2-py3-none-any.whl", hash = "sha256:9170634f126f8efdae22fb58ae8a0eaa86f38365bc57897a6c4f781d1f5875bd", size = 21168, upload-time = "2026-02-16T03:56:08.891Z" }, ] +[[package]] +name = "pluggy" +version = "1.6.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/f9/e2/3e91f31a7d2b083fe6ef3fa267035b518369d9511ffab804f839851d2779/pluggy-1.6.0.tar.gz", hash = "sha256:7dcc130b76258d33b90f61b658791dede3486c3e6bfb003ee5c9bfb396dd22f3", size = 69412, upload-time = "2025-05-15T12:30:07.975Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/54/20/4d324d65cc6d9205fabedc306948156824eb9f0ee1633355a8f7ec5c66bf/pluggy-1.6.0-py3-none-any.whl", hash = "sha256:e920276dd6813095e9377c0bc5566d94c932c33b27a3e3945d8389c374dd4746", size = 20538, upload-time = "2025-05-15T12:30:06.134Z" }, +] + [[package]] name = "post-training" version = "0.1.0" source = { virtual = "." } dependencies = [ + { name = "datasets", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, { name = "huggingface-hub", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, { name = "jinja2", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, { name = "omegaconf", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, @@ -789,13 +806,13 @@ dependencies = [ [package.optional-dependencies] dev = [ { name = "pre-commit", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "pytest", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, ] flash-attn-2 = [ { name = "flash-attn", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, ] trl = [ { name = "accelerate", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, - { name = "datasets", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, { name = "deepspeed", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, { name = "flash-attn", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, { name = "kernels", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, @@ -809,7 +826,7 @@ trl = [ [package.metadata] requires-dist = [ { name = "accelerate", marker = "extra == 'trl'" }, - { name = "datasets", marker = "extra == 'trl'", specifier = ">4.5.0" }, + { name = "datasets", specifier = "==5.0.0" }, { name = "deepspeed", marker = "extra == 'trl'" }, { name = "flash-attn", marker = "extra == 'flash-attn-2'", url = "https://github.com/mjun0812/flash-attention-prebuild-wheels/releases/download/v0.7.16/flash_attn-2.8.3%2Bcu126torch2.9-cp312-cp312-linux_x86_64.whl" }, { name = "flash-attn", marker = "extra == 'trl'", url = "https://github.com/mjun0812/flash-attention-prebuild-wheels/releases/download/v0.7.16/flash_attn-2.8.3%2Bcu126torch2.9-cp312-cp312-linux_x86_64.whl" }, @@ -820,6 +837,7 @@ requires-dist = [ { name = "omegaconf" }, { name = "pre-commit", marker = "extra == 'dev'" }, { name = "psutil" }, + { name = "pytest", marker = "extra == 'dev'" }, { name = "pyyaml" }, { name = "tensorboard", marker = "extra == 'trl'" }, { name = "torch", marker = "extra == 'trl'", specifier = "==2.9.1+cu126" }, @@ -933,6 +951,21 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/c7/21/705964c7812476f378728bdf590ca4b771ec72385c533964653c68e86bdc/pygments-2.19.2-py3-none-any.whl", hash = "sha256:86540386c03d588bb81d44bc3928634ff26449851e99741617ecb9037ee5ec0b", size = 1225217, upload-time = "2025-06-21T13:39:07.939Z" }, ] +[[package]] +name = "pytest" +version = "9.1.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "iniconfig", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "packaging", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "pluggy", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "pygments", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/e4/47/b9efed96c114afcfa3c9d3fe98a76a1d14c74a9e266d397cf6eb64be5e01/pytest-9.1.1.tar.gz", hash = "sha256:1088fbde8f2b49d95a549a195707afa7a76a3ce9bcadc26b6d71f0ffda5fe313", size = 1636369, upload-time = "2026-06-19T10:58:32.857Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/24/25/1de2678b631f5a49215c6c96fff41ba892b0a34df68d6d80292b1b48aa7f/pytest-9.1.1-py3-none-any.whl", hash = "sha256:37a86b45efb9a47a61a36449063e8e18d0cab3161329fc099eb21783169c4f0c", size = 386536, upload-time = "2026-06-19T10:58:31.347Z" }, +] + [[package]] name = "python-dateutil" version = "2.9.0.post0" @@ -1149,7 +1182,7 @@ dependencies = [ { name = "typing-extensions", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, ] wheels = [ - { url = "https://download.pytorch.org/whl/cu126/torch-2.9.1%2Bcu126-cp312-cp312-manylinux_2_28_x86_64.whl", hash = "sha256:67e9b1054f435d33af6fa67343f93d73dc2d37013623672d6ffb24ce39b666c2" }, + { url = "https://download-r2.pytorch.org/whl/cu126/torch-2.9.1%2Bcu126-cp312-cp312-manylinux_2_28_x86_64.whl", hash = "sha256:67e9b1054f435d33af6fa67343f93d73dc2d37013623672d6ffb24ce39b666c2", upload-time = "2026-01-26T16:43:45Z" }, ] [[package]] @@ -1186,7 +1219,7 @@ name = "triton" version = "3.5.1" source = { registry = "https://download.pytorch.org/whl/cu126" } wheels = [ - { url = "https://download.pytorch.org/whl/triton-3.5.1-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:8ce9bcedbd12b393d1309fb899e06818624c7e7765c5231ea5749850015c74d8" }, + { url = "https://download-r2.pytorch.org/whl/triton-3.5.1-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:8ce9bcedbd12b393d1309fb899e06818624c7e7765c5231ea5749850015c74d8", upload-time = "2026-01-22T23:20:53Z" }, ] [[package]] @@ -1235,8 +1268,9 @@ wheels = [ name = "typing-extensions" version = "4.15.0" source = { registry = "https://download.pytorch.org/whl/cu126" } +sdist = { url = "https://files.pythonhosted.org/packages/72/94/1a15dd82efb362ac84269196e94cf00f187f7ed21c242792a923cdb1c61f/typing_extensions-4.15.0.tar.gz", hash = "sha256:0cea48d173cc12fa28ecabc3b837ea3cf6f38c6d1136f85cbaaf598984861466" } wheels = [ - { url = "https://download.pytorch.org/whl/typing_extensions-4.15.0-py3-none-any.whl" }, + { url = "https://files.pythonhosted.org/packages/18/67/36e9267722cc04a6b9f15c7f3441c2363321a3ea07da7ae0c0707beb2a9c/typing_extensions-4.15.0-py3-none-any.whl", hash = "sha256:f0fa19c6845758ab08074a0cfa8b7aecb71c999ca73d62883bc25cc018c4e548" }, ] [[package]] From 9cb46e614e6a09a6a6de3d5b6a0280118a69a9ee Mon Sep 17 00:00:00 2001 From: Arjun Krishnakumar Date: Wed, 16 Sep 2026 15:44:29 +0200 Subject: [PATCH 2/5] docs: add a guide for SFT on a checkpoint with Singularity --- README.md | 162 +++++++++++++++++++++++++++++++++++++++++++++++++++++- 1 file changed, 161 insertions(+), 1 deletion(-) diff --git a/README.md b/README.md index 76c092d..eeb7734 100644 --- a/README.md +++ b/README.md @@ -9,6 +9,7 @@ This repo supports two training backends: ## Table of Contents - [Quick Start](#quick-start) +- [SFT on a Checkpoint with Singularity](#-sft-on-a-checkpoint-with-singularity) - [Project Structure](#-project-structure) - [Design Philosophy](#-design-philosophy) - [Feature Guide](#-feature-guide) @@ -102,6 +103,164 @@ For cluster environments, use the submission script. It auto-generates a SLURM b python scripts/submit.py --config configs/trl/sft.yaml ``` +For the full tokenize-then-train workflow in a container, see [SFT on a Checkpoint with Singularity](#-sft-on-a-checkpoint-with-singularity). + +## 🚀 SFT on a Checkpoint with Singularity + +This guide fine-tunes a given checkpoint with SFT on a SLURM cluster, with training inside a Singularity (or Apptainer) container. It takes two jobs, both submitted from the login node with the same config: + +1. **[Tokenize the datasets](#step-1-tokenize-the-datasets)**: a `--tokenize-only` job on 1 GPU. It loads, filters, tokenizes, and packs the data, writes the result to the Hugging Face datasets cache, and exits. +2. **[Train](#step-2-train)**: the full job. It finds the processed data in the cache, skips preprocessing, and trains. + +Tokenizing first keeps the multi-node allocation from sitting idle during CPU-bound preprocessing, and it surfaces data and chat-template problems in a small job. The examples use paths from LUMI; replace them with your cluster's. + +### Before you start + +#### Install the submission environment + +The login node only prefetches assets and submits jobs; the training stack lives in the container. From the repository root, install the base dependencies: + +```bash +uv sync +source .venv/bin/activate +``` + +Run every `submit.py` command from the repository root. Relative paths in the config (`container.env_file`, `paths.output_base`) resolve against it, and `submit.py` copies the code from it. + +#### Configure the container and the checkpoint + +Start from the reference config: + +```bash +cp configs/trl/sft.yaml configs/trl/my-sft.yaml +``` + +Then set the run name, container, checkpoint, data, and SLURM fields. This example fine-tunes a 9B checkpoint on LUMI: + +```yaml +run_name: oellm-9b-256k-theta64m-prelude-anneal300b-sft # fixed, so both jobs share one run directory +offline: false # true if compute nodes have no internet + +container: + image: /scratch/project_465002530/containers/post-training-rocm7.2.4-py3.12-torch2.9.1-trl1.7.0-olmo-patched.sif + bind_mounts: + - /pfs/lustrep3/scratch/project_465002530/users/krishnak/post-training/ # the repository, with outputs/ + - /scratch/project_465002530/users/krishnak + path: /opt/venv/bin:/opt/rocm/bin:/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin + env_file: env/lumi.env + +model: + name_or_path: birgermoell/oellm-9b-256k-theta64m-prelude-anneal300b # the checkpoint to fine-tune + attn_implementation: flash_attention_2 + dtype: bfloat16 + tokenizer_name_or_path: openeurollm/tokenizer-256k # tokenizer from a separate repo + tokenizer_revision: qwen3-tokens + +training: + max_steps: null # the reference config sets max_steps; keep exactly one duration field + num_train_epochs: 2 + +sft: + max_seq_length: 32768 + packing: true + dataset_num_proc: 32 # workers for tokenizing and packing (null = one process) + +data: + chat_template: qwen3 # must carry {% generation %} markers + num_proc: 32 # workers for loading, transforms, and filtering + datasets: + - name: "dolci-instruct-sft" + path: "allenai/Dolci-Instruct-SFT" + split: "train" + weight: 1.0 + transform: null + +slurm: + account: "project_465002530" + partition: "standard-g" + num_nodes: 4 + gpus_per_node: 8 + cpus_per_task: 56 + mem: "256G" + wall_time: "36:00:00" # keep the quotes +``` + +The container fields: + +- **`container.image`**: the job runs `accelerate launch scripts/train.py` in this image through `singularity exec`. The image must hold the Python packages from `pyproject.toml`, with a PyTorch build for the cluster's GPUs. The `post_training` code does not come from the image; see `run_name` below. +- **`container.path`**: the job sets `PATH` inside the container to exactly this value, so it must contain the directory with `python` and `accelerate`. This image keeps them in `/opt/venv/bin`. The default is `/usr/local/bin:/usr/bin:/bin`. +- **`container.bind_mounts`**: Singularity `--bind` specs. Bind every host path the job reads or writes: the run directory under `paths.output_base`, the Hugging Face cache, and any local checkpoint or dataset. Bind each path as `src` alone, so it keeps the same path inside the container: + - `submit.py` resolves `paths.output_base` to its real path, following symlinks. Bind that real path; here, the repository under `/pfs/lustrep3/...`. + - The frozen config refers to the prefetched checkpoint and tokenizer by their host paths in the Hugging Face cache. +- **`container.env_file`**: a shell file that sets the Hugging Face cache; see [the next section](#write-the-env-file). +- **`run_name`**: at submission, `submit.py` copies `src/post_training/` and `scripts/` into the run directory, and the job runs that copy. With a fixed `run_name`, Step 2 reuses the copy from Step 1, so both jobs run the same transforms and chat templates. The Step 2 submission review warns that the frozen source "will NOT be replaced"; that is expected. To pick up a code change, delete `/src` and `/scripts`, then run Step 1 again. + +#### Write the env file + +The job sources `container.env_file` on the host before it starts the container, then passes the Hugging Face cache variables into the container. The repository ships `env/jupiter.env` as an example. Create one for your cluster, such as `env/lumi.env`: + +```bash +export HF_HOME=/scratch///hf_cache +export HF_HUB_CACHE=$HF_HOME/hub +export HUGGINGFACE_HUB_CACHE=$HF_HOME/hub +export HF_DATASETS_CACHE=$HF_HOME/datasets +``` + +- Export `HF_HOME`, `HF_HUB_CACHE`, and `HUGGINGFACE_HUB_CACHE`. The job script runs with `set -u`, so a missing one stops it with `unbound variable`. `HF_DATASETS_CACHE` defaults to `$HF_HOME/datasets`. +- Use `export NAME=value` lines. `submit.py` reads these lines before it prefetches, so the login node downloads into the cache that the job reads. +- Keep `HF_HOME` inside a bind mount. + +### Step 1: Tokenize the datasets + +```bash +python scripts/submit.py --config configs/trl/my-sft.yaml --tokenize-only +``` + +On the login node, `submit.py`: + +1. reads the Hugging Face cache variables from the env file, +2. downloads the checkpoint, tokenizer, and datasets into that cache (`prefetch_assets: true`, the default), +3. prints a submission review and asks for confirmation (`--confirm` skips it), +4. freezes the config and code into the run directory, and submits the job on 1 node with 1 GPU. The other `slurm` values (account, partition, CPUs, memory, wall time) stay as configured. + +In the container, the job loads the tokenizer and chat template, then loads and filters the datasets. It builds the trainer, which loads the checkpoint, then tokenizes and packs the data. It prints one decoded sample and exits. + +Preprocessing is CPU-bound. Keep `data.num_proc` and `sft.dataset_num_proc` at or below `slurm.cpus_per_task`, and give the job enough wall time. `slurm.*` overrides do not change the processed data, so Step 1 can use its own: + +```bash +python scripts/submit.py --config configs/trl/my-sft.yaml --tokenize-only 'slurm.wall_time="08:00:00"' +``` + +> [!NOTE] +> Quote `slurm.wall_time` on the command line as shown. Unquoted, `24:00:00` parses as the integer `86400`, which SLURM reads as minutes. + +Before Step 2, read `/slurm/slurm-.out`: + +- It shows the `Tokenized dataset preview` block and `--tokenize-only set — exiting after trainer initialization.` Check that the preview follows the chat template's format. +- A warning `... rows, ... with an all-zero assistant mask` means more than 1% of the rows were dropped. A warning that rows are cut `PART-WAY THROUGH their supervised span` means those rows train on truncated answers. Raise `sft.max_seq_length`, or set `sft.truncated_span_action: drop`, then run Step 1 again. +- A `ValueError` stops the job if the chat template lacks `{% generation %}` markers or if no row survives the filter. + +### Step 2: Train + +```bash +python scripts/submit.py --config configs/trl/my-sft.yaml +``` + +Use the same config and the same overrides as Step 1, except for `slurm.*`. `submit.py` renders `/slurm/job.sh` again without `--tokenize-only` and submits it on all nodes. In the container, each preprocessing stage finds its output in the datasets cache and loads it, and training starts. Before the wall time runs out, the job requeues itself and resumes from the latest checkpoint in `/checkpoints/`. + +The cache is hit only when every input to the data pipeline is unchanged. Between the two steps, keep these identical: + +| Keep identical | Why | +|---|---| +| `data.*` | datasets, weights, transforms, seed, and chat template | +| `sft.max_seq_length`, `sft.packing`, `sft.truncated_span_action` | row filtering, truncation, and packing | +| `model.name_or_path`, `model.revision`, `model.tokenizer_name_or_path`, `model.tokenizer_revision` | the tokenizer | +| `container.image` | the library versions that compute the cache keys | +| `container.env_file` | the cache location (`HF_DATASETS_CACHE`) | +| `run_name` | the frozen transforms and chat templates | + +To confirm the cache hit, open the training job's `/slurm/slurm-.err`: the `Tokenizing train dataset` and `Packing train dataset` progress bars must not appear. If they do, an input in the table changed, or `datasets` warned in Step 1 that a function `couldn't be hashed properly`. Either way, the training job processes the data again from scratch. + ## 📂 Project Structure ```text @@ -268,6 +427,7 @@ Templates that are safe for SFT today: |------|--------|-------| | `olmo3-instruct-sft` | `allenai/OLMo-3-7B-Instruct-SFT` (HF Hub) | Use to reproduce the Instruct-SFT recipe. | | `olmo3-think-sft` | `allenai/Olmo-3-7B-Think-SFT` (HF Hub) | Use to reproduce the Think-SFT recipe. | +| `qwen3` | `Qwen/Qwen3-8B` (HF Hub) | Assistant turns whose `` block the template strips stay out of the loss. | Templates that are *not* safe for SFT (kept for inference / DPO compatibility): @@ -333,7 +493,7 @@ You must specify exactly one determining factor for training duration in the `tr - **Debug**: `debug.enabled: true` Forces `report_to: none`, uses a separate output directory, and allows overwriting existing runs. - **Tokenize only**: `--tokenize-only` (CLI flag on `train.py` / `submit.py`) - Exits immediately after the trainer is initialized — dataset loading, tokenization, and packing all run, but the training loop is never entered. Useful for pretokenizing the dataset before committing to a full run. When passed to `submit.py`, the job is automatically constrained to 1 node and 1 GPU. + Exits immediately after the trainer is initialized — dataset loading, tokenization, and packing all run, but the training loop is never entered. Useful for pretokenizing the dataset before committing to a full run. When passed to `submit.py`, the job is automatically constrained to 1 node and 1 GPU. See [SFT on a Checkpoint with Singularity](#-sft-on-a-checkpoint-with-singularity) for the full workflow. ```bash python scripts/submit.py --config configs/trl/sft.yaml --tokenize-only From 407e683ebbbde7e29063a3c5ffe4e661fc9e0de1 Mon Sep 17 00:00:00 2001 From: Arjun Krishnakumar Date: Wed, 16 Sep 2026 15:50:31 +0200 Subject: [PATCH 3/5] feat: add SFT config for the prelude checkpoint on LUMI --- configs/trl/prelude-sft.yaml | 122 +++++++++++++++++++++++++++++++++++ 1 file changed, 122 insertions(+) create mode 100644 configs/trl/prelude-sft.yaml diff --git a/configs/trl/prelude-sft.yaml b/configs/trl/prelude-sft.yaml new file mode 100644 index 0000000..d56e974 --- /dev/null +++ b/configs/trl/prelude-sft.yaml @@ -0,0 +1,122 @@ + +method: sft +backend: trl +run_name: oellm-9b-256k-theta64m-prelude-anneal300b-sft # auto-generated from model + datasets if null +offline: false + +# Container (remove or set image: null for bare-metal) +container: + image: /scratch/project_465002530/containers/post-training-rocm7.2.4-py3.12-torch2.9.1-trl1.7.0-olmo-patched.sif + bind_mounts: + - /pfs/lustrep3/scratch/project_465002530/users/krishnak/post-training/ + - /scratch/project_465002530/users/krishnak + path: /opt/venv/bin:/opt/rocm/bin:/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin + env_file: env/lumi.env + +# -- Model ------------------------------------------------------------------- +model: + name_or_path: birgermoell/oellm-9b-256k-theta64m-prelude-anneal300b + attn_implementation: flash_attention_2 + dtype: bfloat16 + tokenizer_name_or_path: openeurollm/tokenizer-256k + tokenizer_revision: qwen3-tokens + +# -- Training hyper-parameters ----------------------------------------------- +training: + num_train_epochs: 2 + learning_rate: 8.0e-5 + effective_batch_size: 32 # per_device * grad_accum * world_size + per_device_train_batch_size: 1 + warmup_ratio: 0.03 + adam_beta2: 0.95 + lr_scheduler_type: "linear" + gradient_checkpointing: true + bf16: true + seed: 42 + use_liger_kernel: true + +# -- SFT method parameters --------------------------------------------------- +sft: + max_seq_length: 32768 + packing: true + dataset_num_proc: 32 + +# -- Checkpointing ----------------------------------------------------------- +checkpointing: + save_steps: 250 + save_total_limit: 3 # Full checkpoints to keep + inference_checkpoint_steps: 250 # Minimal inference model interval (set to null to disable) + inference_checkpoint_path: "inference_checkpoints" # Relative to run dir + +# -- Data mix ---------------------------------------------------------------- +data: + chat_template: qwen3 # Name from chat template registry + num_proc: 32 # null = auto-detect, capped at 32 + datasets: + - name: "dolci-instruct-sft" + path: "allenai/Dolci-Instruct-SFT" + split: "train" + weight: 1.0 + transform: null # null = already conversational + +# -- DeepSpeed --------------------------------------------------------------- +deepspeed: + bf16: + enabled: auto + zero_optimization: + stage: 2 + overlap_comm: false + contiguous_gradients: true + reduce_scatter: true + gradient_clipping: 1.0 + train_micro_batch_size_per_gpu: "auto" + gradient_accumulation_steps: "auto" + train_batch_size: "auto" + optimizer: + type: AdamW + params: + lr: "auto" + betas: "auto" + eps: "auto" + weight_decay: "auto" + +# -- Accelerate launch flags (explicit multi-node control) ------------------- +accelerate: + mixed_precision: "bf16" + use_deepspeed: true + deepspeed_multinode_launcher: "standard" # "standard" | "pdsh" | etc. + same_network: true # All nodes on same network + rdzv_backend: "static" # "static" | "c10d" | "etcd" + dynamo_backend: "inductor" # "inductor" | "no" | etc. + +# -- Logging & tracking ------------------------------------------------------ +logging: + report_to: + - "wandb" + - "tensorboard" + wandb_project: "sft-training" + logging_steps: 10 + include_num_input_tokens_seen: "non_padding" + +# -- SLURM ------------------------------------------------------------------- +slurm: + account: "project_465002530" + partition: "standard-g" + num_nodes: 4 + gpus_per_node: 8 + cpus_per_task: 56 + wall_time: "36:00:00" + job_name: "Prelude-Context-Extended-Anneal-SFT" + signal_time_seconds: 300 # SIGUSR1 sent this many seconds before timeout to trigger self-healing + max_failures: 1 # Self-healing retry limit + mem: "256G" + +# -- Debug mode -------------------------------------------------------------- +debug: + enabled: false + override_existing: false + +# -- Output paths ------------------------------------------------------------- +paths: + output_base: "outputs" + debug_base: "outputs/debug" From 664ef569de56e94b6e4906c3b799b530621e8f6e Mon Sep 17 00:00:00 2001 From: Arjun Krishnakumar Date: Wed, 16 Sep 2026 15:50:31 +0200 Subject: [PATCH 4/5] docs: point the SFT guide at prelude-sft.yaml --- README.md | 81 +++++++------------------------------------------------ 1 file changed, 9 insertions(+), 72 deletions(-) diff --git a/README.md b/README.md index eeb7734..75b3fa8 100644 --- a/README.md +++ b/README.md @@ -112,80 +112,17 @@ This guide fine-tunes a given checkpoint with SFT on a SLURM cluster, with train 1. **[Tokenize the datasets](#step-1-tokenize-the-datasets)**: a `--tokenize-only` job on 1 GPU. It loads, filters, tokenizes, and packs the data, writes the result to the Hugging Face datasets cache, and exits. 2. **[Train](#step-2-train)**: the full job. It finds the processed data in the cache, skips preprocessing, and trains. -Tokenizing first keeps the multi-node allocation from sitting idle during CPU-bound preprocessing, and it surfaces data and chat-template problems in a small job. The examples use paths from LUMI; replace them with your cluster's. +Tokenizing first keeps the multi-node allocation from sitting idle during CPU-bound preprocessing, and it surfaces data and chat-template problems in a small job. -### Before you start - -#### Install the submission environment - -The login node only prefetches assets and submits jobs; the training stack lives in the container. From the repository root, install the base dependencies: - -```bash -uv sync -source .venv/bin/activate -``` - -Run every `submit.py` command from the repository root. Relative paths in the config (`container.env_file`, `paths.output_base`) resolve against it, and `submit.py` copies the code from it. - -#### Configure the container and the checkpoint - -Start from the reference config: - -```bash -cp configs/trl/sft.yaml configs/trl/my-sft.yaml -``` - -Then set the run name, container, checkpoint, data, and SLURM fields. This example fine-tunes a 9B checkpoint on LUMI: - -```yaml -run_name: oellm-9b-256k-theta64m-prelude-anneal300b-sft # fixed, so both jobs share one run directory -offline: false # true if compute nodes have no internet +The example config is [`configs/trl/prelude-sft.yaml`](configs/trl/prelude-sft.yaml). It fine-tunes a 9B checkpoint on LUMI, with the tokenizer from a separate repo. For your own run, copy it and replace the checkpoint, data, container paths, and SLURM account. -container: - image: /scratch/project_465002530/containers/post-training-rocm7.2.4-py3.12-torch2.9.1-trl1.7.0-olmo-patched.sif - bind_mounts: - - /pfs/lustrep3/scratch/project_465002530/users/krishnak/post-training/ # the repository, with outputs/ - - /scratch/project_465002530/users/krishnak - path: /opt/venv/bin:/opt/rocm/bin:/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin - env_file: env/lumi.env - -model: - name_or_path: birgermoell/oellm-9b-256k-theta64m-prelude-anneal300b # the checkpoint to fine-tune - attn_implementation: flash_attention_2 - dtype: bfloat16 - tokenizer_name_or_path: openeurollm/tokenizer-256k # tokenizer from a separate repo - tokenizer_revision: qwen3-tokens - -training: - max_steps: null # the reference config sets max_steps; keep exactly one duration field - num_train_epochs: 2 - -sft: - max_seq_length: 32768 - packing: true - dataset_num_proc: 32 # workers for tokenizing and packing (null = one process) +### Before you start -data: - chat_template: qwen3 # must carry {% generation %} markers - num_proc: 32 # workers for loading, transforms, and filtering - datasets: - - name: "dolci-instruct-sft" - path: "allenai/Dolci-Instruct-SFT" - split: "train" - weight: 1.0 - transform: null +The login node needs only the base dependencies from [Installation](#installation), because the training stack lives in the container. Run every `submit.py` command from the repository root, inside that environment. Relative paths in the config (`container.env_file`, `paths.output_base`) resolve against the root, and `submit.py` copies the code from it. -slurm: - account: "project_465002530" - partition: "standard-g" - num_nodes: 4 - gpus_per_node: 8 - cpus_per_task: 56 - mem: "256G" - wall_time: "36:00:00" # keep the quotes -``` +#### Configure the container -The container fields: +The fields in `prelude-sft.yaml` that the container run depends on: - **`container.image`**: the job runs `accelerate launch scripts/train.py` in this image through `singularity exec`. The image must hold the Python packages from `pyproject.toml`, with a PyTorch build for the cluster's GPUs. The `post_training` code does not come from the image; see `run_name` below. - **`container.path`**: the job sets `PATH` inside the container to exactly this value, so it must contain the directory with `python` and `accelerate`. This image keeps them in `/opt/venv/bin`. The default is `/usr/local/bin:/usr/bin:/bin`. @@ -213,7 +150,7 @@ export HF_DATASETS_CACHE=$HF_HOME/datasets ### Step 1: Tokenize the datasets ```bash -python scripts/submit.py --config configs/trl/my-sft.yaml --tokenize-only +python scripts/submit.py --config configs/trl/prelude-sft.yaml --tokenize-only ``` On the login node, `submit.py`: @@ -228,7 +165,7 @@ In the container, the job loads the tokenizer and chat template, then loads and Preprocessing is CPU-bound. Keep `data.num_proc` and `sft.dataset_num_proc` at or below `slurm.cpus_per_task`, and give the job enough wall time. `slurm.*` overrides do not change the processed data, so Step 1 can use its own: ```bash -python scripts/submit.py --config configs/trl/my-sft.yaml --tokenize-only 'slurm.wall_time="08:00:00"' +python scripts/submit.py --config configs/trl/prelude-sft.yaml --tokenize-only 'slurm.wall_time="08:00:00"' ``` > [!NOTE] @@ -243,7 +180,7 @@ Before Step 2, read `/slurm/slurm-.out`: ### Step 2: Train ```bash -python scripts/submit.py --config configs/trl/my-sft.yaml +python scripts/submit.py --config configs/trl/prelude-sft.yaml ``` Use the same config and the same overrides as Step 1, except for `slurm.*`. `submit.py` renders `/slurm/job.sh` again without `--tokenize-only` and submits it on all nodes. In the container, each preprocessing stage finds its output in the datasets cache and loads it, and training starts. Before the wall time runs out, the job requeues itself and resumes from the latest checkpoint in `/checkpoints/`. From 2130f7cda94112b5a047e700fb1cfb9dbc241290 Mon Sep 17 00:00:00 2001 From: Arjun Krishnakumar Date: Wed, 16 Sep 2026 16:04:30 +0200 Subject: [PATCH 5/5] docs: add comments to prelude-sft.yaml --- configs/trl/prelude-sft.yaml | 21 ++++++++++----------- 1 file changed, 10 insertions(+), 11 deletions(-) diff --git a/configs/trl/prelude-sft.yaml b/configs/trl/prelude-sft.yaml index d56e974..33e5b2a 100644 --- a/configs/trl/prelude-sft.yaml +++ b/configs/trl/prelude-sft.yaml @@ -8,18 +8,17 @@ offline: false container: image: /scratch/project_465002530/containers/post-training-rocm7.2.4-py3.12-torch2.9.1-trl1.7.0-olmo-patched.sif bind_mounts: - - /pfs/lustrep3/scratch/project_465002530/users/krishnak/post-training/ - - /scratch/project_465002530/users/krishnak + - /pfs/lustrep3/scratch/project_465002530/users/krishnak/post-training/ # replace with your own path to the post-training repo path: /opt/venv/bin:/opt/rocm/bin:/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin - env_file: env/lumi.env + env_file: env/lumi.env # make sure you have this file in the post-training repo, or change to your own env file # -- Model ------------------------------------------------------------------- model: name_or_path: birgermoell/oellm-9b-256k-theta64m-prelude-anneal300b - attn_implementation: flash_attention_2 + attn_implementation: flash_attention_2 # No flash attention 3 for rocm yet, so use flash attention 2 for now dtype: bfloat16 tokenizer_name_or_path: openeurollm/tokenizer-256k - tokenizer_revision: qwen3-tokens + tokenizer_revision: qwen3-tokens # This tokenizer has dedicated tokens for qwen3 chat template, which is used in the data section below. If you use a different tokenizer, make sure to change the chat_template in the data section accordingly. # -- Training hyper-parameters ----------------------------------------------- training: @@ -43,7 +42,7 @@ sft: # -- Checkpointing ----------------------------------------------------------- checkpointing: - save_steps: 250 + save_steps: 250 # How frequently to save the full checkpoints (with optimizer states) save_total_limit: 3 # Full checkpoints to keep inference_checkpoint_steps: 250 # Minimal inference model interval (set to null to disable) inference_checkpoint_path: "inference_checkpoints" # Relative to run dir @@ -54,7 +53,7 @@ data: num_proc: 32 # null = auto-detect, capped at 32 datasets: - name: "dolci-instruct-sft" - path: "allenai/Dolci-Instruct-SFT" + path: "allenai/Dolci-Instruct-SFT" # HuggingFace dataset path split: "train" weight: 1.0 transform: null # null = already conversational @@ -65,7 +64,7 @@ deepspeed: enabled: auto zero_optimization: stage: 2 - overlap_comm: false + overlap_comm: false # Enabling overlap_comm has caused issues before, so we disable it for now. If you want to enable it, set this to true and test carefully. contiguous_gradients: true reduce_scatter: true gradient_clipping: 1.0 @@ -99,14 +98,14 @@ logging: include_num_input_tokens_seen: "non_padding" # -- SLURM ------------------------------------------------------------------- -slurm: +slurm: # LUMI Cluster SLURM configuration account: "project_465002530" - partition: "standard-g" + partition: "standard-g" # Test on dev-g first num_nodes: 4 gpus_per_node: 8 cpus_per_task: 56 wall_time: "36:00:00" - job_name: "Prelude-Context-Extended-Anneal-SFT" + job_name: "Prelude-SFT" signal_time_seconds: 300 # SIGUSR1 sent this many seconds before timeout to trigger self-healing max_failures: 1 # Self-healing retry limit mem: "256G"