[Bugfix] Fix multigpu dispatch_for_generation

[kylesayrs]

Purpose

• Fix test_oneshot_and_finetune_with_tokenizer.py with multiple cuda devices
  • This test had two failures, the first failure occurred as a result of sequential onloading introducing dispatch_for_generation. This function is also used to dispatch for training. However, this method did not account for no split modules
  • The second failure is an existing failure where HFTrainer.compute_loss does not account for multi-gpu models. This will be fixed in the next release by Update Loss Functions to Accept Tensor num_items_in_batch huggingface/transformers#38029

Changes

• Pass no split modules when computing a device map for generation (and training)
• Load model on CPU (since this is now the default flow as of sequential onloading landing

Testing

• Ran test_oneshot_and_finetune_with_tokenizer with two GPUs to completion (with upstream transformers)

[github-actions]

👋 Hi! Thank you for contributing to llm-compressor. Please add the ready label when the PR is ready for review.

Note: This is required to complete the testing suite, please only add the label once the PR is code complete and local testing has been performed.

[gemini-code-assist]

Summary of Changes

Hello @kylesayrs, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request addresses two critical bug fixes primarily affecting model dispatching in multi-GPU setups and the compute_loss function within the HFTrainer. It ensures that models are correctly mapped to devices by properly handling no_split_modules and resolves a type mismatch in loss calculation, improving the robustness and reliability of the training and generation pipelines. Additionally, it updates a key test to reflect current model loading practices.

Highlights

• Bugfix: Device Mapping for Generation: Corrected the dispatch_for_generation utility to properly account for no_split_modules when inferring the device map. This resolves an issue where models with no-split modules were not correctly dispatched across multiple devices, particularly in multi-GPU environments, and ensures consistency between generation and training dispatch logic.
• Bugfix: Loss Computation Type Mismatch: Fixed a type error in HFTrainer.compute_loss where num_items_in_batch was expected as an integer but could be passed as a torch.Tensor. The method now correctly casts the tensor to an integer using .item() before further processing, resolving a runtime failure.
• Test Improvement: Model Loading Strategy: Updated the test_oneshot_and_finetune_with_tokenizer to load models onto the CPU by default by removing the explicit device_map="cuda:0" argument. This aligns the test with the new default model loading flow introduced by sequential onloading, where models are initially loaded to CPU.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point in your pull request via creating an issue comment (i.e. comment on the pull request page) using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in issue comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist is currently in preview and may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments to provide feedback.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[gemini-code-assist]

Code Review

This pull request addresses two bugs related to device dispatching and type handling, and updates the test suite to reflect changes in default model loading behavior. The fixes are targeted and improve the correctness and reliability of the code.

[dsikka]

Can confirm that this fixes the other two nightly errors as well:

tests/llmcompressor/transformers/finetune/test_oneshot_and_finetune.py
and tests/llmcompressor/transformers/finetune/test_finetune_no_recipe_custom_dataset.py

[rahul-tuli]

Blazing fast resolution!

[kylesayrs]

@dsikka I've confirmed that these tests run until the original loss = loss / num_items_in_batch issue, which is fixed in upstream transformers