Struct MegatronInstaller 

pub struct MegatronInstaller { /* private fields */ }

The Megatron-LM installer.

Implementations

impl MegatronInstaller

pub fn new(config: MegatronConfig) -> Self

Create a new Megatron installer with the given config.

pub fn with_defaults() -> Self

Create with default config.

pub fn dependencies(&self) -> &[&str]

Get the list of required dependencies.

Megatron-LM depends on:

• PyTorch (ROCm-enabled)
• MPI4Py

pub fn rocm_env(&self, rocm_env: &RocmEnv) -> Vec<(String, String)>

Get the ROCm environment variables needed for Megatron.

The original script sets:

• AMD_LOG_LEVEL=0
• HIP_VISIBLE_DEVICES
• ROCR_VISIBLE_DEVICES
• HSA_OVERRIDE_GFX_VERSION=11.0.0
• PYTORCH_ROCM_ARCH=gfx1100
• ROCM_PATH
• HSA_TOOLS_LIB
• PYTORCH_ALLOC_CONF=expandable_segments:True

pub fn build_git_clone_command(&self) -> ShellCommand

Construct the git clone command for Megatron-LM.

The original script clones: git clone https://github.com/NVIDIA/Megatron-LM.git <dir>

It retries up to 3 times on failure.

pub fn max_clone_retries(&self) -> u32

Get the maximum number of clone retries.

pub fn build_pip_install_command(&self) -> ShellCommand

Construct the pip install command for Megatron-LM.

The original script installs with: pip install -e . --no-deps

The editable install is done from the cloned directory. Dependencies are resolved separately (safe requirements filtering).

pub fn build_dep_install_command(&self, package: &str) -> ShellCommand

Construct the pip install command for a specific dependency.

Used for installing individual safe dependencies.

pub fn build_requirements_install_command( &self, requirements_path: &str, ) -> ShellCommand

Construct a pip install command for a requirements file.

pub fn is_safe_package(&self, package_name: &str) -> bool

Check if a Python package is safe to install (non-CUDA/non-NVIDIA).

The original script filters out packages that would conflict with the ROCm environment: nvidia-, cuda, torch, torchvision, etc.

pub fn check_build_output( &self, stdout: &str, stderr: &str, exit_code: i32, ) -> Result<(), String>

Check build output for errors.