search

caikit / caikit-nlp

Apache License 2.0
12 stars 39 forks source link

Enable FSDP for Prompt Tuning via Peft #177

Open alex-jw-brooks opened 10 months ago

alex-jw-brooks commented 10 months ago

Description

As a user of prompt tuning, I want to be able to leverage multiple GPUs at train time!

Discussion

Extends https://github.com/caikit/caikit-nlp/issues/175 to leverage PyTorch FSDP; there may be some refactoring to consolidate with fine tuning FSDP configuration / dataset stuff, which may be part of this story, or left to a follow up. We will decide once we see how Peft / trainer integration goes.

Acceptance Criteria

chakrn commented 8 months ago

@gkumbhat can you also link the relevant PR to this one?