[docker, CI] use radixark/Megatron-LM' and allow specify sglang/megatron PR in CI

[yueming-yuan]

ci-megatron-pr: #11
ci-sglang-pr: sglang-miles

[gemini-code-assist]

Summary of Changes

Hello @yueming-yuan, 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 refactors the Dockerfile responsible for setting up the development environment. The primary goal is to move away from maintaining and applying a custom patch for Megatron-LM, instead opting for direct integration with a specific branch of an external repository. This change simplifies dependency management by utilizing pre-built wheels and standardizing the installation of various deep learning libraries, leading to a more robust and easier-to-maintain build process.

Highlights

• Megatron-LM Integration Strategy: The integration of Megatron-LM has been updated to directly use the miles-main branch from the radixark/Megatron-LM repository, replacing the previous method of applying a local megatron.patch file.
• Streamlined Dependency Management: The Dockerfile now leverages pre-built Python wheels from a specified GitHub repository (yueming-yuan/miles-wheels), simplifying the installation process for several dependencies like flash-attn, TransformerEngine, and apex.
• Improved CUDA 13 Support: Conditional logic for CUDA 13 installations has been refined, particularly for TransformerEngine and nvidia-cudnn, ensuring more straightforward and version-specific dependency handling.
• Dockerfile Clean-up: Removed outdated patch application steps and conditional build logic, resulting in a cleaner and more maintainable Dockerfile.

Changelog

• docker/Dockerfile.dev
  • Added comments detailing different image build configurations for radixark/miles.
  • Updated build arguments to specify MEGATRON_REPO, MEGATRON_BRANCH, and WHEELS_REPO.
  • Introduced a new step to download pre-built Python wheels from the WHEELS_REPO.
  • Simplified installation commands for flash-attn, flash_attn_3, TransformerEngine, and apex to use the newly downloaded pre-built wheels.
  • Modified the Megatron-LM cloning process to use the specified MEGATRON_REPO and MEGATRON_BRANCH.
  • Removed the explicit megatron.patch application logic.
  • Removed conditional Triton and sgl-kernel installations.
  • Adjusted nvidia-cudnn installation to be conditional on ENABLE_CUDA_13.
• docker/patch/dev/megatron.patch
  • Removed the megatron.patch file, as its functionality is now integrated via direct repository usage.

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 significantly simplifies the Dockerfile by replacing the megatron.patch with a direct dependency on a forked radixark/Megatron-LM repository. This is a great improvement for maintainability. The changes also refactor the installation of Python dependencies, using pre-built wheels more effectively and cleaning up the logic for different CUDA versions. The new method for downloading all wheels from a GitHub release in one step is particularly clever. I have one suggestion to improve the readability and robustness of the Python script used for downloading the wheels.

[gemini-code-assist]

This Python one-liner is a bit dense and has a couple of potential issues:

1. It will fail with a KeyError if the GitHub API response for the release does not contain an assets key. Using .get('assets', []) would be more robust.
2. Using a list comprehension for its side effects (calling subprocess.run) is not idiomatic Python. A for loop is more explicit and readable.

Consider refactoring this into a multi-line script within the RUN command for better readability and maintainability, for example:

import sys, json, subprocess

release_data = json.load(sys.stdin)
for asset in release_data.get('assets', []):
    if asset['name'].endswith('.whl'):
        url = asset['browser_download_url']
        filename = '/tmp/wheels/' + asset['name']
        print(f"Downloading {url} to {filename}")
        subprocess.run(['curl', '-fSL', '-o', filename, url], check=True)