DOCSEARCHbyAlgolia
Submit Docs

torch_musa Documentation

repository·main·Indexed 19 days ago

https://github.com/moorethreads/torch_musa

An extended Python package for PyTorch that enables hardware acceleration on Moore Threads MUSA graphics cards. It provides API consistency with PyTorch and CUDA, supporting MUSA-specific modules, distributed training via the mccl backend, and tools for operator micro-benchmarking. The package includes utilities for CUDA-porting, structured kernel implementation strategies (Legacy, LegacyMeta, and Customized), and C++ deployment demos for TorchScript models.

Tokens
49.9K
Snippets
135
Records
169
Agent score
65%

What's inside torch_musa

  1. Overview of torch_musa

    main

    What is torch_musa?

    torch_musa is a Python package that extends PyTorch to enable users to leverage the computing power of Moore Threads (MUSA) GPUs.

    Key Features

    • API Compatibility: The interface is consistent with PyTorch. Users can migrate smoothly by switching device strings from cpu or cuda to musa.
    • Migration Tools: Provides tools for CUDA code migration, building MUSA extensions, and debugging.
    • Advanced Optimizations: Supports features like Dynamic Double Casting and Unified Memory Management.
    • C/C++ Extension API: Offers an efficient API for writing custom operators or network layers in C/C++ with minimal boilerplate.

  2. Overview of the Operator Comparison and Tracing Tool

    main

    The Operator Comparison and Tracing tool is a suite of debugging utilities designed to ensure the accuracy and reliability of PyTorch model operators on MUSA hardware. It provides several key capabilities for developers and researchers:

    • CPU Operator Comparison: Verifies custom or MUSA-specific operator implementations by comparing their outputs against standard CPU results to ensure consistency across hardware platforms.
    • Module Tracing: Provides context and hierarchy for operators, allowing developers to trace anomalies back to specific parts of the model structure.
    • NaN/Inf Detection: Uses the NanInfTracker functionality to locate numerical instabilities (NaNs and Infs) directly on the GPU with minimal performance impact.
    • Anomaly Handling: Offers strategies to isolate, analyze, and resolve detected issues, such as isolating specific operators, adjusting tolerance levels, or using whitelists for expected discrepancies.
    • Training Step Control: Allows for adaptive activation of debugging features (e.g., during AMP training) to avoid misleading NaN/Inf values during initial training phases.
    • Distributed Support: Designed for multi-GPU environments, supporting log management and selective activation on specific ranks.

  3. What is the musa_converter tool?

    main

    The musa_converter is a one-click conversion tool provided by torch_musa developers. It is designed to automatically transform native PyTorch training or inference scripts originally written for CUDA platforms into scripts that can run directly on the MUSA platform.

    Warning: The conversion process modifies your source files directly. It is highly recommended to back up your code before running the tool.

  4. What is torch_musa?

    main

    Overview

    torch_musa is a plugin developed by Moore Threads to enable support for Moore Threads GPUs within the PyTorch framework. It is built on top of PyTorch v2.0.0 and is designed to be decoupled from the core PyTorch code to facilitate easier maintenance and upgrades.

    Key Features

    • Plugin Architecture: Uses PyTorch's third-party backend extension interfaces to dynamically register Moore Threads high-performance computing libraries.
    • CUDA Compatibility: Includes a CUDA compatibility module that allows PyTorch community CUDA kernels to run on Moore Threads GPUs. The CUDA porting process is automated during the torch_musa compilation process, reducing the cost of operator adaptation.
    • Low Migration Cost: The Python frontend interfaces are kept consistent with the PyTorch community's CUDA interfaces, minimizing the learning curve and the effort required to migrate existing models.

  5. What is torch_musa and how does it work

    main

    Overview

    torch_musa is a plugin developed by Moore Threads to enable support for Moore Threads GPUs within the PyTorch framework. It is built on top of PyTorch v2.0.0 and is designed to be decoupled from the core PyTorch code to facilitate easier maintenance and upgrades.

    Key Features

    • Plugin Architecture: It uses PyTorch's third-party backend extension interfaces to dynamically register Moore Threads high-performance computing libraries into PyTorch.
    • CUDA Compatibility: It includes a CUDA compatibility module. This allows CUDA kernels from the PyTorch community to run on Moore Threads GPUs after porting. Crucially, much of the CUDA porting work is performed automatically during the torch_musa compilation process, reducing the cost of operator adaptation.
    • Consistent API: The Python frontend interfaces are kept consistent with the PyTorch community's CUDA interfaces, minimizing the learning curve and migration effort for users moving from CUDA-based workflows.

    Core Components

    torch_musa is organized into several key directories:

    • torch_musa/core: Provides the Python frontend interfaces for modules such as amp, device, memory, stream, and event.
    • torch_musa/csrc: Contains the C++ implementation code:
      • csrc/amp: C++ implementation for mixed-precision modules.
      • csrc/aten: C++ Tensor library, including MUDNN operator adaptation and CUDA-Porting operator adaptation.
      • csrc/core: Core functional libraries including device management, memory allocation management, Stream management, and Events management.
      • csrc/distributed: C++ implementation for distributed modules.

  6. Implement Structured Operators in MUSA

    main

    For C++ operators with multiple calling rules (functional, inplace, out) within a structured group, PyTorch uses a multi-level inheritance pattern. To implement a structured operator in MUSA, you must define the relationship between these rules in native_functions.yaml so that torchgen can automatically generate the necessary class hierarchies and registration code.

    Inheritance Hierarchy (Bottom to Top)

    1. meta::structured_{name}: Inherits from MetaBase. Implements the meta function for preprocessing.
    2. native::structured_{name}_{backend}: Inherits from the meta class. Implements the impl function for computation.
    3. structured_{name}_{backend}_functional: Inherits from the impl class. Overrides methods for result tensor creation.
    4. structured_{name}_{backend}_inplace: Inherits from the impl class. Overrides methods for in-place tensor validation.
    5. structured_{name}_{backend}_out: Inherits from the impl class. Overrides methods for result tensor validation and size changes.

    Configuration in native_functions.yaml

    To ensure correct codegen, use the following fields:

    • structured_delegate: Used in functional and inplace definitions to point to the out operator (the 'reduction' target).
    • structured: True: Used in the out operator definition to mark it as structured.
    • dispatch: Specifies the backend implementation (e.g., CUDA: tril_cuda or PrivateUse1: MusaTril).
    • structured_inherits: (Optional) Specifies a custom base class instead of the default MetaBase.
    • precomputed: (Optional) Defines the mapping of intermediate variables passed from meta to impl functions.
    - func: tril(Tensor self, int diagonal=0) -> Tensor
  structured_delegate: tril.out

- func: tril_(Tensor(a!) self, int diagonal=0) -> Tensor(a!)
  structured_delegate: tril.out

- func: tril.out(Tensor self, int diagonal=0, *, Tensor(a!) out) -> Tensor(a!)
  structured: True
  dispatch: CUDA: tril_cuda

  7. How to use device guards in MUSA operators

    main

    When implementing custom operators in C++, you must use c10::musa::MUSAGuard or c10::musa::OptionalMUSAGuard to set the correct device ID at the operator entry point. This ensures that subsequent calls to MUSA-specific libraries (like muDNN) operate on the correct device context.

    When to use a guard: Use a guard if your operator depends on the MUSA device context (e.g., it calls GetMudnnHandle()).

    When NOT to use a guard: If your operator does not depend on the MUSA device context and only uses the device information from existing tensors (e.g., calling .to(device)), a guard is not necessary.

    Tensor NativeDropoutBackward(const Tensor& grad_output, const Tensor& mask, double scale) {
  // Set the device guard using the input tensor's device
  c10::musa::MUSAGuard device_guard(input.device()); 
  
  // Now it is safe to call MUSA-specific handles
  muHandle& h = GetMudnnHandle();
  ...
}

  8. Control debugging activation with start_step and end_step

    main

    In scenarios like Automatic Mixed Precision (AMP) where initial training steps have large scales and frequently produce NaN/Inf, you can control when the debugging tools activate.

    Set start_step and end_step in CompareWithCPU or NanInfTracker. The tools will only be active when start_step <= step_cnt < end_step. You must manually increment the step counter using the .step() method on the context manager object.

    from torch_musa.utils.compare_tool import CompareWithCPU

model = get_your_model()
# Tools only active between step 5 and step 10
with CompareWithCPU(atol=0.001, rtol=0.001, verbose=True, start_step=5, end_step=10) as compare_with_cpu:
    for epoch in range(epoch_num):
        for step in range(step_num):
            train_step(model)
            compare_with_cpu.step()  # Increment the internal step counter

  9. Understand the torch_musa core directory structure

    main

    The torch_musa repository is organized into several key directories that separate Python interfaces from C++ implementations:

    • torch_musa/core: Contains the primary Python modules, providing frontend interfaces for amp, device, memory, stream, and event.
    • torch_musa/csrc: Contains the C++ implementation code:
      • csrc/amp: C++ implementation of mixed-precision modules.
      • csrc/aten: C++ Tensor library, including MUDNN operator adaptation and CUDA-Porting operator adaptation.
      • csrc/core: Core functional libraries, including device management, memory allocation management, Stream management, and Events management.
      • csrc/distributed: C++ implementation of distributed modules.
    • torch_musa/tests: Test files.

  10. How MUSA operator development works with torchgen

    main

    PyTorch separates operator definition from implementation. Definitions (format, implementation method, backend binding, and export rules) are stored in YAML files. For MUSA, torch_musa extends the torchgen logic via a custom codegen module.

    To develop a MUSA operator:

    1. Implement the MUSA-specific calculation logic in C++.
    2. Add the operator's description to torch_musa/csrc/aten/ops/musa_functions.yaml.
    3. During compilation, the codegen module parses this YAML to automatically generate the operator's interface (*.h) and definition (*.cpp) files, which handle the binding to the MUSA backend.

    torch_musa reuses the PrivateUse1 key to implement operator registration, following official PyTorch recommendations.

  11. Reuse Public Functions (Legacy Implementation)

    main

    If an operator's implementation logic is identical across backends (e.g., it only involves view transformations without actual data computation), you can use a 'Legacy' implementation.

    By declaring the MUSA dispatch name to match the existing CPU/CUDA function name in musa_functions.yaml, the codegen module will recognize the match and automatically call the common function during registration. This avoids redundant code.

    Example for view_as_real:

    # MUSA declaration matching CPU/CUDA to reuse the common implementation
- func: view_as_real
  dispatch:    PrivateUse1: view_as_real

  12. Scaled-dot-product Attention Computing in Torch Musa

    main

    Torch Musa implements torch.nn.functional.scaled_dot_product_attention with two available computing modes, mirroring PyTorch CUDA behavior:

    1. Math mode: A general-purpose mode suitable for all types of attention computing, though it may not be the most performance-optimized.
    2. FlashAttention mode: An optimized mode that supports dtype=torch.float16 with head dimensions that are less than 64, less than 128, or equal to 160.

    Warning for Training: FlashAttention mode currently does not support backward passes. If you use FlashAttention during the forward pass, the backward pass will fail. For training workloads, you must disable FlashAttention.

    import torch
import torch.nn.functional as F

# Example setup for attention tensors
query = torch.rand(32, 8, 128, 64, dtype=torch.float16, device="cuda")
key = torch.rand(32, 8, 128, 64, dtype=torch.float16, device="cuda")
value = torch.rand(32, 8, 128, 64, dtype=torch.float16, device="cuda")