Skip to content

[tunix/sft] Add shard_input_data flag to PeftTrainer#1325

Open
copybara-service[bot] wants to merge 1 commit intomainfrom
test_891043878
Open

[tunix/sft] Add shard_input_data flag to PeftTrainer#1325
copybara-service[bot] wants to merge 1 commit intomainfrom
test_891043878

Conversation

@copybara-service
Copy link
Copy Markdown

[tunix/sft] Add shard_input_data flag to PeftTrainer

Adds the shard_input_data boolean flag to TrainingConfig (defaulting to True) to control whether inputs should be automatically sharded within the training and evaluation loops.

This provides flexibility for users to disable auto-sharding in cases where the input data is already appropriately sharded before being passed to the trainer (e.g. via an external data loader).

Also adds test_shard_input_data_flag to peft_trainer_test.py to verify that sharding_utils.shard_input is called only when the flag is set to True.

Adds the `shard_input_data` boolean flag to `TrainingConfig` (defaulting to True) to control whether inputs should be automatically sharded within the training and evaluation loops.

This provides flexibility for users to disable auto-sharding in cases where the input data is already appropriately sharded before being passed to the trainer (e.g. via an external data loader).

Also adds `test_shard_input_data_flag` to `peft_trainer_test.py` to verify that `sharding_utils.shard_input` is called only when the flag is set to True.

PiperOrigin-RevId: 891043878
@google-cla
Copy link
Copy Markdown

google-cla bot commented Mar 28, 2026

Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

View this failed invocation of the CLA check for more information.

For the most up to date status, view the checks section at the bottom of the pull request.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant