Currently prompt tuning module only support training on single gpu. In order to support larger models, or larger batch sizes or allow larger text sizes, we would need to provide support to handle training across multiple GPUs (on a single node).