Refactoring AssistedCandidateGenerator for Improved Modularity and Reusability

[keyboardAnt]

What does this PR do?

This PR refactors the AssistedCandidateGenerator and AssistedCandidateGeneratorDifferentTokenizers classes to improve code readability, maintainability, and reusability. By breaking down large methods into smaller helper functions and reusing common logic, we enhance the modularity of these classes. This refactoring lays a cleaner foundation for upcoming features like UniversalSpeculativeDecodingGenerator without introducing any new functionality.

Background

While working on Universal Speculative Decoding (#34760), which introduces the UniversalSpeculativeDecodingGenerator, we identified opportunities to refactor existing code. The goal is to reuse core logic across different candidate generators and simplify the integration of new features that enable speculative decoding across models with different tokenizers.

By submitting this refactoring as a separate PR, we aim to:

• Streamline the review process: Focus on structural improvements without the complexity of new features.
• Accelerate availability: Make these improvements accessible to the community sooner.
• Facilitate future development: Provide a robust base for upcoming enhancements in the generation capabilities.

This refactor is a collaboration with @jmamou, who has already reviewed it (keyboardAnt#1).

Key Changes

1. Code Restructuring

• Decomposed Large Methods: Broke down the get_candidates methods in both classes into smaller, focused helper functions.
  • Example helper functions:
    • _calculate_new_tokens
    • _update_past_and_masks
    • _prepare_generation_args
    • _generate_candidates
• Simplified Initialization: Streamlined the __init__ methods to remove redundancy and enhance clarity.

2. Improved Reusability

• Modularized Operations: Encapsulated unique functionalities in dedicated methods within each class.
  • For AssistedCandidateGeneratorDifferentTokenizers, methods like _prepare_assistant_input_ids and _process_assistant_outputs handle tokenizer-specific logic.
• Reused Common Logic: Leveraged inheritance and method overriding to share code between classes and reduce duplication.

3. Enhanced Readability

• Clear Naming Conventions: Renamed variables and methods for better clarity and understanding.
• Consistent Formatting: Applied uniform code formatting and style guidelines across the classes.
• Documentation: Added docstrings and comments to explain the purpose and functionality of methods.

Motivation

This refactoring is motivated by the need to:

• Reduce Technical Debt: Clean up the codebase to make it more maintainable.
• Facilitate Future Features: Prepare the existing classes for integration with upcoming speculative decoding functionalities.
• Encourage Community Contributions: Make the code more approachable for contributors by improving readability and structure.

By isolating these changes, we enable reviewers to focus solely on the structural improvements without the added complexity of new features. This approach helps maintain a high code quality standard and simplifies the review and merging process.

Dependencies

• No New Dependencies: This PR does not introduce any new dependencies. It solely refactors existing code.

Before Submitting

• Part of Speculative Decoding Effort: This PR is a preparatory step for enhancing speculative decoding capabilities ([OLD] New PR: #35029. [[Universal Speculative Decoding CandidateGenerator]] #34760).
• Contributor Guidelines: I have read and followed the contributor guidelines.
• No New Functionality: This PR strictly involves refactoring without adding new features.
• Testing: All existing tests pass successfully. No new tests are needed as functionality remains unchanged.

Who Can Review?

The following reviewers are well-suited to review this PR: @gante, @ArthurZucker

---

This PR aims to strengthen the foundation for speculative decoding and other future enhancements by improving the existing code's structure and maintainability. We appreciate your time and look forward to your feedback.

[jmamou]

@zucchini-nlp feel free to review it :-)

[zucchini-nlp]

Thanks a lot for making the code more composable! LGTM as nothing changed in terms of functionality. Just left one comment as we had several PRs in parallel modifying the assisted generation code :)

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[zucchini-nlp]

Thanks @keyboardAnt ! LGTM as nothing was changed except for composing code into smaller functions.

The only question is the use of self.prev_target_ids which I've mentioned was removed in prev PR explaining it is not needed. This is what I got from reviewing the prev PR

hmm, so to make sure, that means the prev impl when we checked prev_target_ids was not really correct? And we should check the length of already accepted input_ids

yes

So do we need to save prev token ids from target model or we can re-use the current token ids, because the current token ids in any case will have the prev token ids as prefix with new accepted tokens appended at the end

[jmamou]

Thanks @keyboardAnt ! LGTM as nothing was changed except for composing code into smaller functions.

The only question is the use of self.prev_target_ids which I've mentioned was removed in prev PR explaining it is not needed. This is what I got from reviewing the prev PR

hmm, so to make sure, that means the prev impl when we checked prev_target_ids was not really correct? And we should check the length of already accepted input_ids

yes

So do we need to save prev token ids from target model or we can re-use the current token ids, because the current token ids in any case will have the prev token ids as prefix with new accepted tokens appended at the end

For UAG, we need just the number of tokens in self.prev_target_ids (prev_target_ids.shape[1]) in order to know which target tokens were added after the target validation of the last speculative iteration and to convert them to the relevant draft tokens (instead of converting from the beginning).
I propose to use similar approach as in USD when we store only the previous number of target tokens _prev_target_seq_len.

[zucchini-nlp]

we need just the number of tokens in self.prev_target_ids (prev_target_ids.shape[1]

Ah that makes sense if we don't care about the actual token ids used previously, because the tokens should be available without storing them. Then we can indeed store only the prev length

[keyboardAnt]

we need just the number of tokens in self.prev_target_ids (prev_target_ids.shape[1]

Ah that makes sense if we don't care about the actual token ids used previously, because the tokens should be available without storing them. Then we can indeed store only the prev length

Thanks, @zucchini-nlp. I've replaced self.prev_target_ids with self.prev_target_ids_len and rebased main to keep the history clean. Can you please approve the awaiting workflows?

[jmamou]

@keyboardAnt
please apply the same fix as at keyboardAnt#4.
else SD will not work when target and assistant are not on the same device

[jmamou]

@keyboardAnt please apply the same fix as at keyboardAnt#4. else SD will not work when target and assistant are not on the same device

@keyboardAnt
PR #35116 fixes multi-gpu issue

[zucchini-nlp]

You can now request review from the core maintainer to merge this :)

[jmamou]

You can now request review from the core maintainer to merge this :)

@ArthurZucker

[ArthurZucker]

Looks great! Thanks for the detailed explanations 🤗
Merging!

Need to dismiss it to merge!

[ArthurZucker]

Ah can you resolve conflicts? 🤗

[jmamou]

Ah can you resolve conflicts? 🤗

@ArthurZucker done!

[danielkorat]

@ArthurZucker @zucchini-nlp
conflicts were resolved 😃 ready for merge

[zucchini-nlp]

thanks, merging