KEP-2401: Complement torch plugin to support torchtune config mutation #2587

Electronic-Waste · 2025-04-08T14:15:44Z

What this PR does / why we need it:

This PR adds the torchtune config mutation implementation in torch plugin.

As we discussed before, we'll implement the config mutation/validation in server-side, so as to avoid frequent SDK changes and provide better backward compatibility for users.

In details, this PR:

Add config mutation/mapping for the torchtune configs passed in .spec.trainer.args
Add valiations for torchtune config
Add more UTs for torch EnforceMLPolicy and Validate function

REF: https://github.com/kubeflow/trainer/tree/master/docs/proposals/2401-llm-trainer-v2#complement-torch-plugin

/cc @kubeflow/wg-training-leads @astefanutti @franciscojavierarceo @saileshd1402 @deepanker13 @akshaychitneni

Which issue(s) this PR fixes (optional, in Fixes #<issue number>, #<issue number>, ... format, will close the issue(s) when PR gets merged):
Fixes #2507 #2508

Checklist:

Docs included if any changes are user facing

… plugin. Signed-off-by: Electronic-Waste <[email protected]>

Signed-off-by: Electronic-Waste <[email protected]>

google-oss-prow · 2025-04-08T14:15:53Z

@Electronic-Waste: GitHub didn't allow me to request PR reviews from the following users: saileshd1402.

Note that only kubeflow members and repo collaborators can review this PR, and authors cannot review their own PRs.

In response to this:

What this PR does / why we need it:

This PR adds the torchtune config mutation implementation in torch plugin.

/cc @kubeflow/wg-training-leads @astefanutti @franciscojavierarceo @saileshd1402 @deepanker13 @akshaychitneni

Which issue(s) this PR fixes (optional, in Fixes #<issue number>, #<issue number>, ... format, will close the issue(s) when PR gets merged):
Fixes #2507

Checklist:

Docs included if any changes are user facing

Instructions for interacting with me using PR comments are available here. If you have questions or suggestions related to my behavior, please file an issue against the kubernetes/test-infra repository.

Signed-off-by: Electronic-Waste <[email protected]>

… Args. Signed-off-by: Electronic-Waste <[email protected]>

Signed-off-by: Electronic-Waste <[email protected]>

coveralls · 2025-04-09T08:36:05Z

Pull Request Test Coverage Report for Build 14692575862

Warning: This coverage report may be inaccurate.

This pull request's base commit is no longer the HEAD commit of its target branch. This means it includes changes from outside the original pull request, including, potentially, unrelated coverage changes.

For more information on this, see Tracking coverage changes with pull request builds.
To avoid this issue with future PRs, see these Recommended CI Configurations.
For a quick fix, rebase this PR at GitHub. Your next report should be accurate.

Details

59 of 61 (96.72%) changed or added relevant lines in 1 file are covered.
62 unchanged lines in 1 file lost coverage.
Overall coverage increased (+0.7%) to 67.18%

Changes Missing Coverage	Covered Lines	Changed/Added Lines	%
pkg/runtime/framework/plugins/torch/torch.go	59	61	96.72%

Files with Coverage Reduction	New Missed Lines	%
pkg/runtime/framework/plugins/coscheduling/coscheduling.go	62	0.0%

Totals
Change from base Build 14341999020:	0.7%
Covered Lines:	1789
Relevant Lines:	2663

💛 - Coveralls

Signed-off-by: Electronic-Waste <[email protected]>

Electronic-Waste · 2025-04-09T09:51:40Z

PTAL if you have time, thanks:)

/assign @kubeflow/wg-training-leads @astefanutti @akshaychitneni @franciscojavierarceo @deepanker13

pkg/runtime/framework/plugins/torch/torch.go

sdk/kubeflow/trainer/utils/utils.py

Signed-off-by: Electronic-Waste <[email protected]>

pkg/runtime/framework/plugins/torch/torch.go

Signed-off-by: Electronic-Waste <[email protected]>

Electronic-Waste · 2025-04-27T13:45:47Z

@andreyvelich Thanks for your detailed review! I've addressed your comments. PTAL if you have time:)

andreyvelich

Overall lgtm.
Just small comment.
/lgtm
/assign @tenzen-y @kubeflow/wg-training-leads @saileshd1402 @astefanutti @franciscojavierarceo for the review.

pkg/runtime/framework/plugins/torch/torch.go

andreyvelich

I think, we can move forward.
Let's address any additional changes in the followup PRs.
Thank you for this @Electronic-Waste!
/approve

google-oss-prow · 2025-04-29T13:20:45Z

[APPROVALNOTIFIER] This PR is APPROVED

This pull-request has been approved by: andreyvelich

The full list of commands accepted by this bot can be found here.

The pull request process is described here

Needs approval from an approver in each of these files:

~~OWNERS~~ [andreyvelich]

Approvers can indicate their approval by writing /approve in a comment
Approvers can cancel approval by writing /approve cancel in a comment

Electronic-Waste · 2025-04-29T13:29:28Z

@andreyvelich @astefanutti Thanks for your detailed review! I'll create another issue to discuss about: #2587 (comment).

kubeflow/trainer#2587) * chore(plugin): Add torchtune-related constants & update current torch plugin. Signed-off-by: Electronic-Waste <[email protected]> * chore(plugin): Add EnforceMLPolicy for torchtune. Signed-off-by: Electronic-Waste <[email protected]> * chore(plugin): Add UTs in torch plugin. Signed-off-by: Electronic-Waste <[email protected]> * fix(test): fix error in torch plugin UTs. Signed-off-by: Electronic-Waste <[email protected]> * chore(plugin): Choose recipe according to numNodes & numProcPerNode & Args. Signed-off-by: Electronic-Waste <[email protected]> * chore(sdk): Add PretrainedModel enum type. Signed-off-by: Electronic-Waste <[email protected]> * chore(plugin): Add torchtune config arg. Signed-off-by: Electronic-Waste <[email protected]> * chore(test): add UT for single-device full fine-tuning with torchtune. Signed-off-by: Electronic-Waste <[email protected]> * chore(test): Add test for multi-nodes full fine-tuning with torchtune. Signed-off-by: Electronic-Waste <[email protected]> * chore(test): Update torch validate UTs. Signed-off-by: Electronic-Waste <[email protected]> * fix(lint): fix lint error. Signed-off-by: Electronic-Waste <[email protected]> * fix(sdk): remove pretrained model enum type in sdk. Signed-off-by: Electronic-Waste <[email protected]> * fix(plugin): retrieve model name from runtimeRef. Signed-off-by: Electronic-Waste <[email protected]> * fix(lint): fix typo. Signed-off-by: Electronic-Waste <[email protected]> * fix(plugin): make some adjustments according to the review. Signed-off-by: Electronic-Waste <[email protected]> * fix(sdk): remove runtime in get_trainer_crd_from_builtin_trainer. Signed-off-by: Electronic-Waste <[email protected]> * fix(plugin): pass PET_ env variables in torch plugin for torchtune. Signed-off-by: Electronic-Waste <[email protected]> * fix(plugin): add env validation for torchtune. Signed-off-by: Electronic-Waste <[email protected]> * fix(plugin): update comments. Signed-off-by: Electronic-Waste <[email protected]> * fix(plugins): fix the implementation according to the review. Signed-off-by: Electronic-Waste <[email protected]> * test(plugins): fix UT error in torch plugin. Signed-off-by: Electronic-Waste <[email protected]> * fix: fix UT and e2e tests error. Signed-off-by: Electronic-Waste <[email protected]> * fix: remove debug info. Signed-off-by: Electronic-Waste <[email protected]> * fix(test): add args in UTs related to torchtune. Signed-off-by: Electronic-Waste <[email protected]> * fix(test): update torchtune related args. Signed-off-by: Electronic-Waste <[email protected]> * fix(test): Add a UT for multi-node mode check in torch plugin. Signed-off-by: Electronic-Waste <[email protected]> --------- Signed-off-by: Electronic-Waste <[email protected]>

kubeflow#2587) * chore(plugin): Add torchtune-related constants & update current torch plugin. Signed-off-by: Electronic-Waste <[email protected]> * chore(plugin): Add EnforceMLPolicy for torchtune. Signed-off-by: Electronic-Waste <[email protected]> * chore(plugin): Add UTs in torch plugin. Signed-off-by: Electronic-Waste <[email protected]> * fix(test): fix error in torch plugin UTs. Signed-off-by: Electronic-Waste <[email protected]> * chore(plugin): Choose recipe according to numNodes & numProcPerNode & Args. Signed-off-by: Electronic-Waste <[email protected]> * chore(sdk): Add PretrainedModel enum type. Signed-off-by: Electronic-Waste <[email protected]> * chore(plugin): Add torchtune config arg. Signed-off-by: Electronic-Waste <[email protected]> * chore(test): add UT for single-device full fine-tuning with torchtune. Signed-off-by: Electronic-Waste <[email protected]> * chore(test): Add test for multi-nodes full fine-tuning with torchtune. Signed-off-by: Electronic-Waste <[email protected]> * chore(test): Update torch validate UTs. Signed-off-by: Electronic-Waste <[email protected]> * fix(lint): fix lint error. Signed-off-by: Electronic-Waste <[email protected]> * fix(sdk): remove pretrained model enum type in sdk. Signed-off-by: Electronic-Waste <[email protected]> * fix(plugin): retrieve model name from runtimeRef. Signed-off-by: Electronic-Waste <[email protected]> * fix(lint): fix typo. Signed-off-by: Electronic-Waste <[email protected]> * fix(plugin): make some adjustments according to the review. Signed-off-by: Electronic-Waste <[email protected]> * fix(sdk): remove runtime in get_trainer_crd_from_builtin_trainer. Signed-off-by: Electronic-Waste <[email protected]> * fix(plugin): pass PET_ env variables in torch plugin for torchtune. Signed-off-by: Electronic-Waste <[email protected]> * fix(plugin): add env validation for torchtune. Signed-off-by: Electronic-Waste <[email protected]> * fix(plugin): update comments. Signed-off-by: Electronic-Waste <[email protected]> * fix(plugins): fix the implementation according to the review. Signed-off-by: Electronic-Waste <[email protected]> * test(plugins): fix UT error in torch plugin. Signed-off-by: Electronic-Waste <[email protected]> * fix: fix UT and e2e tests error. Signed-off-by: Electronic-Waste <[email protected]> * fix: remove debug info. Signed-off-by: Electronic-Waste <[email protected]> * fix(test): add args in UTs related to torchtune. Signed-off-by: Electronic-Waste <[email protected]> * fix(test): update torchtune related args. Signed-off-by: Electronic-Waste <[email protected]> * fix(test): Add a UT for multi-node mode check in torch plugin. Signed-off-by: Electronic-Waste <[email protected]> --------- Signed-off-by: Electronic-Waste <[email protected]>

Electronic-Waste added 3 commits April 8, 2025 11:47

chore(plugin): Add torchtune-related constants & update current torch…

7babe4e

… plugin. Signed-off-by: Electronic-Waste <[email protected]>

chore(plugin): Add EnforceMLPolicy for torchtune.

e6a9c54

Signed-off-by: Electronic-Waste <[email protected]>

chore(plugin): Add UTs in torch plugin.

afa49c4

Signed-off-by: Electronic-Waste <[email protected]>

google-oss-prow bot requested review from deepanker13, akshaychitneni, a team, astefanutti and franciscojavierarceo April 8, 2025 14:15

google-oss-prow bot added the size/L label Apr 8, 2025

Electronic-Waste marked this pull request as draft April 8, 2025 14:16

google-oss-prow bot added the do-not-merge/work-in-progress label Apr 8, 2025

Electronic-Waste changed the title ~~KEP-2401: Complement torch plugin to support torchtune config mutation~~ [WIP] KEP-2401: Complement torch plugin to support torchtune config mutation Apr 8, 2025

Electronic-Waste added 4 commits April 8, 2025 15:59

fix(test): fix error in torch plugin UTs.

ada86f8

Signed-off-by: Electronic-Waste <[email protected]>

chore(plugin): Choose recipe according to numNodes & numProcPerNode &…

c9e340e

… Args. Signed-off-by: Electronic-Waste <[email protected]>

chore(sdk): Add PretrainedModel enum type.

8ecaa5b

Signed-off-by: Electronic-Waste <[email protected]>

chore(plugin): Add torchtune config arg.

06b555d

Signed-off-by: Electronic-Waste <[email protected]>

Electronic-Waste force-pushed the feat/torchtune-plugin branch from 4865ba0 to 06b555d Compare April 9, 2025 06:20

Electronic-Waste added 2 commits April 9, 2025 08:23

chore(test): add UT for single-device full fine-tuning with torchtune.

ba55d4c

Signed-off-by: Electronic-Waste <[email protected]>

chore(test): Add test for multi-nodes full fine-tuning with torchtune.

206822e

Signed-off-by: Electronic-Waste <[email protected]>

chore(test): Update torch validate UTs.

a29dddc

Signed-off-by: Electronic-Waste <[email protected]>

google-oss-prow bot added size/XL and removed size/L labels Apr 9, 2025

fix(lint): fix lint error.

a9f993a

Signed-off-by: Electronic-Waste <[email protected]>

Electronic-Waste marked this pull request as ready for review April 9, 2025 09:39

Electronic-Waste changed the title ~~[WIP] KEP-2401: Complement torch plugin to support torchtune config mutation~~ KEP-2401: Complement torch plugin to support torchtune config mutation Apr 9, 2025

google-oss-prow bot removed the do-not-merge/work-in-progress label Apr 9, 2025

Electronic-Waste commented Apr 9, 2025

View reviewed changes

pkg/runtime/framework/plugins/torch/torch.go Outdated Show resolved Hide resolved

sdk/kubeflow/trainer/utils/utils.py Outdated Show resolved Hide resolved

google-oss-prow bot added size/L and removed size/XL labels Apr 22, 2025

test(plugins): fix UT error in torch plugin.

79816eb

Signed-off-by: Electronic-Waste <[email protected]>

andreyvelich reviewed Apr 25, 2025

View reviewed changes

pkg/runtime/framework/plugins/torch/torch.go Outdated Show resolved Hide resolved

fix: fix UT and e2e tests error.

71b4b5b

Signed-off-by: Electronic-Waste <[email protected]>

google-oss-prow bot added size/XL and removed size/L labels Apr 27, 2025

Electronic-Waste added 4 commits April 27, 2025 12:21

fix: remove debug info.

47321b7

Signed-off-by: Electronic-Waste <[email protected]>

fix(test): add args in UTs related to torchtune.

40f4119

Signed-off-by: Electronic-Waste <[email protected]>

fix(test): update torchtune related args.

52287a9

Signed-off-by: Electronic-Waste <[email protected]>

fix(test): Add a UT for multi-node mode check in torch plugin.

87f68ff

Signed-off-by: Electronic-Waste <[email protected]>

andreyvelich reviewed Apr 28, 2025

View reviewed changes

pkg/runtime/framework/plugins/torch/torch.go Show resolved Hide resolved

google-oss-prow bot assigned andreyvelich Apr 28, 2025

google-oss-prow bot added the lgtm label Apr 28, 2025

andreyvelich reviewed Apr 29, 2025

View reviewed changes

google-oss-prow bot added the approved label Apr 29, 2025

google-oss-prow bot merged commit 040b34e into kubeflow:master Apr 29, 2025
20 checks passed

google-oss-prow bot added this to the v2.0 milestone Apr 29, 2025

Electronic-Waste deleted the feat/torchtune-plugin branch April 29, 2025 13:26

This was referenced Apr 29, 2025

KEP-2401: Validate fine-tuning configurations in torch plugin #2508

Closed

Retrieve arguments from Trainer instance kubeflow/sdk#19

Open

Electronic-Waste mentioned this pull request Jun 14, 2025

fix(manifests): Update manifests to enable LLM fine-tuning workflow with CTR and TrainJob yaml files #2669

Merged

1 task

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

KEP-2401: Complement torch plugin to support torchtune config mutation #2587

KEP-2401: Complement torch plugin to support torchtune config mutation #2587

Uh oh!

Electronic-Waste commented Apr 8, 2025 •

edited

Loading

Uh oh!

google-oss-prow bot commented Apr 8, 2025

Uh oh!

coveralls commented Apr 9, 2025 •

edited

Loading

Uh oh!

Electronic-Waste commented Apr 9, 2025

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Electronic-Waste commented Apr 27, 2025

Uh oh!

andreyvelich left a comment

Uh oh!

Uh oh!

andreyvelich left a comment

Uh oh!

google-oss-prow bot commented Apr 29, 2025

Uh oh!

Uh oh!

Electronic-Waste commented Apr 29, 2025

Uh oh!

Uh oh!

KEP-2401: Complement torch plugin to support torchtune config mutation #2587

KEP-2401: Complement torch plugin to support torchtune config mutation #2587

Uh oh!

Conversation

Electronic-Waste commented Apr 8, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

google-oss-prow bot commented Apr 8, 2025

Uh oh!

coveralls commented Apr 9, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Pull Request Test Coverage Report for Build 14692575862

Warning: This coverage report may be inaccurate.

Details

💛 - Coveralls

Uh oh!

Electronic-Waste commented Apr 9, 2025

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Electronic-Waste commented Apr 27, 2025

Uh oh!

andreyvelich left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

andreyvelich left a comment

Choose a reason for hiding this comment

Uh oh!

google-oss-prow bot commented Apr 29, 2025

Uh oh!

Uh oh!

Electronic-Waste commented Apr 29, 2025

Uh oh!

Uh oh!

Electronic-Waste commented Apr 8, 2025 •

edited

Loading

coveralls commented Apr 9, 2025 •

edited

Loading