Skip to content

Add support for Custom Accelerate Instance in Trainer#40633

Open
Flakes342 wants to merge 5 commits into
huggingface:mainfrom
Flakes342:feature/cust-acc
Open

Add support for Custom Accelerate Instance in Trainer#40633
Flakes342 wants to merge 5 commits into
huggingface:mainfrom
Flakes342:feature/cust-acc

Conversation

@Flakes342

Copy link
Copy Markdown
Contributor

What does this PR do?

Fixes #40616

This PR adds the ability to pass a custom Accelerator instance to the Trainer constructor, giving users fine-grained control over distributed training configurations while maintaining all existing Trainer functionality.

Before submitting

Description

Currently, users who need advanced distributed training configurations (custom DeepSpeed settings, specific FSDP configurations, etc.) have limited options to customize the Accelerator used by Trainer. The existing approach forces users to either modify TrainingArguments which doesn't expose all accelerate options to them or use accelerator_config which is less intuitive and limited.

This change allows users to create and configure their own Accelerator instance with full control over:

  • DeepSpeed configurations with custom plugins
  • FSDP settings and policies
  • Tensor parallelism configurations
  • Custom dataloader behaviors
  • Advanced gradient accumulation strategies
  • Experiment tracking and logging setup

Changes

  • Modified Trainer initialization to add an optional accelerator parameter
  • Added postprocess_accelerator method: Handles postprocessing for user-provided accelerators - similar to current create_accelerator_and_postprocess method
  • Refactored accelerator initialization: Clean separation between creation and postprocessing logic

Example

Adding a simple example of how custom accelerator functionality can be used in trainer:

from accelerate import Accelerator
from accelerate.utils import DeepSpeedPlugin

# Custom DeepSpeed configuration
deepspeed_plugin = DeepSpeedPlugin(
    zero_stage=3,
    offload_optimizer_device="cpu",
    offload_param_device="cpu",
    zero3_init_flag=True,
    zero3_save_16bit_model=True
)

accelerator = Accelerator(deepspeed_plugin=deepspeed_plugin)

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=dataset,
    accelerator=accelerator
)

Testing

Added comprehensive tests to trainer/test_trainer.py covering:

  • Basic custom accelerator functionality
  • Backward compatibility (no accelerator parameter)
  • Postprocessing attribute setup
  • Settings preservation
  • Gather function configuration
  • End-to-end training integration

Please let me know if the tests need to be written in some other way. Also, please review and provide your comments and I'd be happy to incorporate any corrections/suggestions.

Who can review? @Rocketknight1, @SunMarc, @anyone :)

raise ValueError("save_only_model option is not compatible with FSDP state dict type 'SHARDED_STATE_DICT'")


def postprocess_accelerator(self):

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I wanted to create this as a different method to support any modifications to custom accelerate function in future too.

@Flakes342

Copy link
Copy Markdown
Contributor Author

Again, I am not sure what these test failures are..

@Flakes342

Copy link
Copy Markdown
Contributor Author

Changes for this PR have been done in:
src/transformers/trainer.py
tests/trainer/tests_trainer.py

Rest all seems to have happened because of auto formatting fix

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.

Allow users to pass through Accelerate instance

1 participant