Tell Us: What Would Make Trainer Better?

[SunMarc]

RFC: Trainer improvements

The Trainer class is a core component of the Transformers library, and we're looking to make it even better.
We're gathering inputs on potential improvements, new features, and pain points you've experienced with the Trainer class.

We're particularly interested in feedback on:

• Training Performance: Speed optimizations, memory efficiency, distributed training improvements
• API & Usability: API design, documentation, ease of use, common pain points
• Features & Functionality: Missing features, callbacks, logging, evaluation, checkpointing
• Integration: Better integration with other libraries or tools
• Debugging & Monitoring: Better error messages, debugging tools, training diagnostics
• Developer Experience: Contributing to Trainer, extending functionality, custom trainers

💭 Questions to Guide Your Feedback

For Users

• What's the most frustrating aspect of using Trainer?
• What feature would save you the most time?
• Are there any workarounds you regularly implement that should be built-in?
• What use cases does Trainer not support well?
• How could the documentation be improved?

For Developers

• What would make it easier to extend or customize Trainer?
• Are there architectural changes that would improve maintainability?
• What integration points with other libraries should we prioritize?
• Where do you see technical debt or areas needing refactoring?

📝 How to Contribute Your Ideas

Please share your thoughts by:

1. Commenting below with your feature requests or improvement ideas
2. Describing your use case - help us understand the problem you're trying to solve
3. Being specific - concrete examples and scenarios are extremely helpful
4. Upvoting ideas from others that you'd also find valuable

🤝 What Happens Next

• We'll review all suggestions and feedback
• Popular and high-impact ideas will be prioritized
• Contributors interested in implementing features are welcome to volunteer!

🙏 Thank You

Your feedback drives the evolution of the Transformers library. Every perspective helps us build better tools for the ML community. Special thanks to @stas00 for getting this conversation started !

List of PRs created thanks to this thread

Surface-Level Unbloat

• Simplify TrainingArguments docstring #43568 : better TrainingArguments docstring to easily navigate between the different arguments
• [Trainer] Move sort and rotate checkpoints to standalone functions #43736 : move sort and rotate checkpoints to standalone functions
• [Trainer] Move NEFTune impl to standalone functions #43714 : move netftune to standalone functions
• [Trainer] Move optimizer cls init to trainer_optimizer.py #43738: move optimizers initialization to trainer_optimizers.py file
• Minor changes trainer #43744: bunch of minor changes in Trainer

Methods

Initialization

• Refactor trainer init #43807: Trainer init
• Update TrainingArguments #43806: TrainingArguments init + test

Reroder

• Reorder Trainer methods #43914: Reroder all Trainer methods

Refactor tests

• Refactor trainer data_collator and callbacks tests #43776: data collator and callbacks
• Update common tests Trainer #44260: other common trainer tests
• Update distributed tests #44338: Distributed tests
• Add testing guide for agents for trainer tests #44328: Template for AGENTS when it comes to testing

Priority List

• soon

[chschroeder]

Excellent initiative, thank you!

I am both a user and a developer, but primarily a downstream developer where I maintain small-text, which uses setfit, which in turn uses sentence-transformers, which uses transformers. So this is an integration use case, where I often juggle around different combinations of these libraries, thereby encountering a lot of different errors.

For Users

• Optional functionality should not cause errors: Optional features should never throw errors if unused. Dependency upgrades are always an adventure, and I regularly get annoyed whenever something fails just because I have, for example, not the right huggingface-hub version (despite not using any hub functionality). A well-placed check or slightly adjusted import location could completely prevent such errors.

  Since huggingface-hub might not be seen as optional from HF's point of view, another example would be wandb. Not every user will be interested in every integration.

For Developers

• Complex methods and GPU memory: In my active learning use scenarios, I have a scenario of repeated training within the same process, i.e. long-running programs with multiple Trainer calls. Occasionally (and seemingly random), I suffer from GPU out-of-memory situations, which don't make sense at first glance: GPU-memory-wise, it should not be a problem because the same batch size had been working for many hours (in the same as well in previous processes) before, and moreover, those errors occur only sometimes. Upon closer investigation I found that this is often just unlucky timing of the garbage collection, which prevents GPU memory from being freed, often owed to overly long and complex methods in the trainer. Smaller methods would decrease the local variable context and make it possible to free references to GPU tensors.

  Just a few days ago, as so often I had to override parts of the trainer adding gc.collect()/ torch.cuda.empty_cache() calls as well as explicit del-statement to reduce my out-of-memory problems. This is a common hack I have been using over the years, which should not be needed. From a developer's perspective, breaking things into smaller, clear methods is rarely a downside. My use case might be niche, but long running applications will probably not be that uncommon.

Edit: Slightly clarified my wording to make it more readable.

[ParagEkbote]

I would like to offer a couple of thoughts to help improve documentation and the developer experience of the Trainer class. Firstly, there could be more explicit text in the docs which describes which major open-source libraries the trainer currently supports. For example, I was unaware that there was optuna support was available until found the original HPO page.

Secondly, since Quantization aware training (QAT) has become extremely common in model releases, the existing documentation does not specify what packages beyond bitsandbytes are available for usage while training. It could be better if there is clarity about which methods the community could incorporate to improve depth for model fine-tuning/training.

Lastly, after the release of the kernels library, is there any upcoming plans to add support to integrate it with Trainer? It could further help in using FA2, FA3 or vllm kernels to improve the speed of training and simplify environment setup.

Many thanks to all of the maintainers who have built the trainer class and continue to improve it for benefit of the HF community.

[FremyCompany]

Not sure how easy that would be, but when training a custom model, it would be nice if the config_xxx.py and modeling_xxx.py could be auto-saved (assuming it is possible to specify the files as @model_files(__file__) or something like that (also in the checkpoints).\

I know you said one request but I have another one: right now the trainer accepts a loss, but I would sometimes like to return more than one loss, which should be logged separately and backward()ed consecutively.

[cynricfu]

It would be great if the Trainer API could offer better support for iterable datasets, which are essential for large-scale training.

Related unresolved issues include but not limited to the followings:
#31441
#33544
#34205