Add logging to tiled mlp and fix target chunk size calculation

[mmathew23]

Fix a shard size calculation issue when tiled mlp is set to target. Also add logging so we know Tiled MLP is enabled.

[gemini-code-assist]

Summary of Changes

Hello @mmathew23, 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 enhances the TiledMLP functionality by resolving a critical shard size calculation error that could occur under specific configurations. Additionally, it integrates a new logging mechanism to provide clear indications when TiledMLP is actively being used, improving debugging and operational transparency.

Highlights

• TiledMLP Logging: Introduced new global variables (FIRST_PASS, UNSLOTH_ENABLE_LOGGING, UNSLOTH_ENABLE_TILED_LOGGING) to control and enable logging for TiledMLP, providing visibility when it's active.
• Shard Size Calculation Fix: Corrected the qlen_chunk_size calculation within the TiledMLP.forward method, especially when max_flat_qlen is specified, to prevent issues with shard sizing.
• Parameter Propagation: Modified the TiledMLP.forward method signature to accept max_flat_qlen and updated its usage in tiled_forward_target_gb and tiled_forward_arctic_size to ensure correct parameter passing.
• Backward Method Alignment: Adjusted the return signature of the TiledMLP.backward method to align with the updated forward method, accommodating the new max_flat_qlen parameter.

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 by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

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 pull request 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 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. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

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 introduces logging to the tiled MLP functionality and corrects an issue with shard size calculation. The logic for the fix appears sound and is well-integrated with the existing tiling strategies. The new logging will be helpful for users to know when Tiled MLP is active. I've included a couple of suggestions to align the new logging code with the existing logging utilities in the project, which will improve consistency and maintainability.

[gemini-code-assist]

To avoid code duplication and potential inconsistencies, you should import UNSLOTH_ENABLE_LOGGING from unsloth_zoo.log instead of redefining it here. The version in unsloth_zoo.log is also more robust as it checks for "1", "True", and "true".

Please add from unsloth_zoo.log import UNSLOTH_ENABLE_LOGGING to your imports. You can then replace these two lines with just the definition for UNSLOTH_ENABLE_TILED_LOGGING.

Suggested change

	UNSLOTH_ENABLE_LOGGING = os.environ.get("UNSLOTH_ENABLE_LOGGING", "0") == "1"
	UNSLOTH_ENABLE_TILED_LOGGING = UNSLOTH_ENABLE_LOGGING and os.environ.get("UNSLOTH_ENABLE_TILED_LOGGING", "0") == "1"
	UNSLOTH_ENABLE_TILED_LOGGING = UNSLOTH_ENABLE_LOGGING and os.environ.get("UNSLOTH_ENABLE_TILED_LOGGING", "0") == "1"

[gemini-code-assist]

It's better to use a proper logger for logging instead of print(). This allows for better control over log levels, formatting, and output streams. The unsloth_zoo.log module provides a configured logger instance that you can import and use.

After importing logger from unsloth_zoo.log, you can make the following change:

Suggested change

	print(f"Unsloth: Enabling TiledMLP to reduce VRAM usage! chunk size: {split_sizes[0]}")
	logger.info(f"Unsloth: Enabling TiledMLP to reduce VRAM usage! chunk size: {split_sizes[0]}")