### Multi-Machine Multi-GPU Distributed Setup

Source: https://github.com/moorethreads/torch_musa/blob/main/docs/pdf/developer_guide/source/9_optimize_performance/optimize_performance.md

Provides example usage instructions for setting up distributed training across multiple machines and GPUs. This snippet outlines the command-line arguments and environment variables needed for configuration.

```bash
# Usage:
# e.g:
# On machine A:
#   MASTER_ADDR=xxx MASTER_PORT=xxx python3 resnet50_distributed_ddp_profiler.py -n 2 -g 1 -nr 0
#
# On machine B:

```

--------------------------------

### Project and Environment Setup

Source: https://github.com/moorethreads/torch_musa/blob/main/examples/cpp/kernel/CMakeLists.txt

Initializes the CMake project and sets environment variables for MUSA. This is a standard starting point for CMake projects.

```cmake
cmake_minimum_required(VERSION 3.10 FATAL_ERROR)
project(w8a8 CXX C)
add_definitions(-DHAVE_AVX2_CPU_DEFINITION)
add_definitions(-DHAVE_AVX512_CPU_DEFINITION)
set(ENV{MUSA_HOME} "/usr/local/musa")
```

--------------------------------

### Run C++ Model Deployment Example

Source: https://github.com/moorethreads/torch_musa/blob/main/examples/cpp/README.md

Generate the model and then run the C++ example application with a specified model file. The output shows the model's predictions.

```python
python model_gen.py 
```

```bash
./build/example-app "resnet50.pt"
# -2.0971 -0.9347 -0.3347 -0.6047 -0.9165
# [ PrivateUse1FloatType{1,5} ]
```

--------------------------------

### Install C++ Extension for PyTorch

Source: https://github.com/moorethreads/torch_musa/blob/main/benchmark/operator_benchmark/README.md

Install the required cpp_extension for PyTorch to run TorchMusa operator benchmarks. Navigate to the extension directory and run the setup script.

```bash
cd $PYTORCH_REPO/benchmarks/operator_benchmark/pt_extension
python setup.py install
```

--------------------------------

### Install Dependencies for PDF Generation

Source: https://github.com/moorethreads/torch_musa/blob/main/docs/pdf/README.md

Installs necessary LaTeX packages and project requirements. Ensure you have apt and pip available.

```bash
apt-get update
apt-get install texlive-xetex
apt-get install latexmk
apt-get install latex-cjk-all
apt-get install fonts-noto-cjk

pip install -r requirements.txt
```

--------------------------------

### Using torch.profiler with torch_musa

Source: https://github.com/moorethreads/torch_musa/blob/main/docs/markdown/source/07_optimization.md

This example demonstrates the recommended way to import torch_musa before PyTorch profiler modules for accurate performance analysis. It includes model setup, data loading, and the profiler configuration with a custom trace handler.

```python
import torch
import torch_musa
import torchvision.models as models
from torch.profiler import profile, record_function, schedule
import torch.nn as nn
import torch.optim
import torch.utils.data
import torchvision
import torchvision.transforms as T


model = models.resnet50(weights=models.ResNet50_Weights.IMAGENET1K_V1)
model.musa()
transform = T.Compose([T.Resize(256), T.CenterCrop(224), T.ToTensor()])
trainset = torchvision.datasets.CIFAR10(root='./data', train=True,
                                        download=True, transform=transform)
trainloader = torch.utils.data.DataLoader(trainset, batch_size=32,
                                          shuffle=True, num_workers=0)
criterion = nn.CrossEntropyLoss().musa()
optimizer = torch.optim.SGD(model.parameters(), lr=0.001, momentum=0.9)
device = torch.device("musa:0")
model.train()

my_schedule = schedule(
  skip_first=1,
  wait=1,
  warmup=1,
  active=2,
  repeat=2)

def trace_handler(p):
  output = p.key_averages().table(sort_by="self_musa_time_total", row_limit=10)
  print(output)
  p.export_chrome_trace("/tmp/trace_" + str(p.step_num) + ".json")

# with_stack=True requires experimental_config's setting at torch 2.0.0 which has not resolved the issue, https://github.com/pytorch/pytorch/issues/100253
with profile(
  schedule=my_schedule,
  on_trace_ready=trace_handler,
  profile_memory=True,
  record_shapes=True,
  with_stack=True,
  experimental_config=torch._C._profiler._ExperimentalConfig(verbose=True)
  ) as prof:
  with record_function("model_training"):
    for step, data in enumerate(trainloader, 0):
        print("step:{}".format(step))
        inputs, labels = data[0].to(device=device), data[1].to(device=device)
        outputs = model(inputs)
        loss = criterion(outputs, labels)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        prof.step()
        if step >= 10:
            break

print(prof.key_averages(group_by_input_shape=True).table(sort_by="musa_time_total", row_limit=10))
print(prof.key_averages(group_by_stack_n=5).table(sort_by="self_musa_time_total", row_limit=2))
```

--------------------------------

### Configure and Install Musa Library

Source: https://github.com/moorethreads/torch_musa/blob/main/torch_musa/csrc/extension/CMakeLists.txt

Sets the installation directory for the Musa library if not already defined and then installs the target library. This ensures the compiled library is placed in the correct location.

```cmake
if(NOT EXT_MUSA_KERNEL_INSTALL_LIB_DIR)
  set(EXT_MUSA_KERNEL_INSTALL_LIB_DIR lib)
endif()
install(TARGETS ${EXT_MUSA_KERNEL_LIB}
        DESTINATION ${EXT_MUSA_KERNEL_INSTALL_LIB_DIR})
```

--------------------------------

### Compile torch_musa (Manual Step)

Source: https://github.com/moorethreads/torch_musa/blob/main/docs/markdown/source/03_compile_and_install.md

Installs torch_musa from source. Supports debug and ASan modes.

```shell
cd torch_musa
pip install -r requirements.txt
python setup.py install
# debug mode: DEBUG=1 python setup.py install
# asan mode: USE_ASAN=1 python setup.py install
```

--------------------------------

### Compile PyTorch (Manual Step)

Source: https://github.com/moorethreads/torch_musa/blob/main/docs/markdown/source/03_compile_and_install.md

Installs PyTorch from source after patching. Supports debug and ASan modes.

```shell
cd pytorch
pip install -r requirements.txt
python setup.py install
# debug mode: DEBUG=1 python setup.py install
# asan mode: USE_ASAN=1 python setup.py install
```

--------------------------------

### Clone Torch-MUSA Repository

Source: https://github.com/moorethreads/torch_musa/blob/main/README.md

Clone the torch_musa repository to begin the installation process from source.

```bash
git clone https://github.com/MooreThreads/torch_musa.git
cd torch_musa
```

--------------------------------

### Distributed Training with Torch-MUSA

Source: https://github.com/moorethreads/torch_musa/blob/main/docs/pdf/developer_guide/source/5_quickguide/quickguide.md

Example of setting up and running distributed training using `torch.nn.parallel.DistributedDataParallel` with the 'mccl' backend. Ensure `MASTER_ADDR` and `MASTER_PORT` are correctly set.

```python
"""Demo of DistributedDataParall"""
import os
import torch
from torch import nn
from torch import optim
from torch.nn.parallel import DistributedDataParallel as DDP
import torch.distributed as dist
import torch.multiprocessing as mp
import torch_musa


class Model(nn.Module):
    def __init__(self):
        super().__init__()
        self.linear = nn.Linear(5,5)
    def forward(self, x):
        return self.linear(x)

def start(rank, world_size):
    if os.getenv("MASTER_ADDR") is None:
        os.environ["MASTER_ADDR"]= ip # IP must be specified here
    if os.getenv("MASTER_PORT") is None:
        os.environ["MASTER_PORT"]= port # port must be specified here
    dist.init_process_group("mccl", rank=rank, world_size=world_size)

def clean():
    dist.destroy_process_group()

def runner(rank, world_size):
    torch_musa.set_device(rank)
    start(rank, world_size)
    model = Model().to('musa')
    ddp_model = DDP(model, device_ids=[rank])
    optimizer = optim.SGD(ddp_model.parameters(), lr=0.001)
    for _ in range(5):
        input_tensor = torch.randn(5, dtype=torch.float, requires_grad=True).to('musa')
        target_tensor = torch.zeros(5, dtype=torch.float).to('musa')
        output_tensor = ddp_model(input_tensor)
        loss_f = nn.MSELoss()
        loss = loss_f(output_tensor, target_tensor)
        loss.backward()
        optimizer.step()
    clean()

if __name__ == "__main__":
    mp.spawn(runner, args=(2,), nprocs=2, join=True)
```

--------------------------------

### Install MUSA torchaudio

Source: https://github.com/moorethreads/torch_musa/blob/main/README.md

Install torchaudio from the PyTorch audio source repository using the release branch corresponding to your PyTorch version.

```shell
git clone https://github.com/pytorch/audio.git -b release/${version} --depth 1
cd audio && python setup.py install
```

--------------------------------

### C++ Deployment Example

Source: https://github.com/moorethreads/torch_musa/blob/main/docs/markdown/source/04_quick_start.md

This C++ code demonstrates how to load and run a PyTorch model using torch_musa. Ensure the model is saved using torch jit.trace or jit.script and register the 'musa' backend for PrivateUse1.

```cpp
#include <torch/script.h>
#include <torch_musa/csrc/core/Device.h>
#include <iostream>
#include <memory>

int main(int argc, const char* argv[]) {
    // Register 'musa' for PrivateUse1 as we save model with 'musa'.
    c10::register_privateuse1_backend("musa");

    torch::jit::script::Module module;
    // Load model which saved with torch jit.trace or jit.script.
    module = torch::jit::load(argv[1]);

    std::vector<torch::jit::IValue> inputs;
    // Ready for input data.
    torch::Tensor input = torch::rand({1, 3, 224, 224}).to("musa");
    inputs.push_back(input);

    // Model execute.
    at::Tensor output = module.forward(inputs).toTensor();

    return 0;
}
```

--------------------------------

### Example CSV Output from Benchmark

Source: https://github.com/moorethreads/torch_musa/blob/main/benchmark/operator_benchmark/README.md

This is an example of the CSV output generated by the visualization tool. It includes columns for operator name, mode, test case details, configuration, and various performance metrics like latency and throughput.

```csv
OP               ,Mode  ,TestCase                                                                  ,TestConfig                                                                    ,backward ,macs ,tflops ,Latency ,latency_variance ,latency_mean ,0%      ,25%     ,50%     ,75%     ,100%
RMSNormBenchmark ,Eager ,"RMSNormBenchmark_input_shape(2,128,1024)_musa_dtypetorch.float16"        ,"{\'input_shape\': [2, 128, 1024], \'device\': \'musa\', \'dtype\': \'torch.float16\'}" ,False    ,-1   ,     0 , 29.118 ,           0.069 ,      29.219 , 29.042 , 29.072 , 29.118 , 29.166 , 29.946
RMSNormBenchmark ,Eager ,"RMSNormBenchmark_input_shape(8,32,2048)_musa_dtypetorch.float16"         ,"{\'input_shape\': [8, 32, 2048], \'device\': \'musa\', \'dtype\': \'torch.float16\'}"  ,False    ,-1   ,     0 , 33.784 ,           0.065 ,      33.912 , 33.736 , 33.758 , 33.784 , 33.924 , 34.588
RMSNormBenchmark ,Eager ,"RMSNormBenchmark_input_shape(8,128,2048)_musa_dtypetorch.float16"        ,"{\'input_shape\': [8, 128, 2048], \'device\': \'musa\', \'dtype\': \'torch.float16\'}" ,False    ,-1   ,     0 , 47.877 ,           0.041 ,      47.962 , 47.829 , 47.84  , 47.877 , 48.018 , 48.493
RMSNormBenchmark ,Eager ,"RMSNormBenchmark_input_shape(1,128,4096)_musa_dtypetorch.float16"        ,"{\'input_shape\': [1, 128, 4096], \'device\': \'musa\', \'dtype\': \'torch.float16\'}" ,False    ,-1   ,     0 , 41.593 ,           0.047 ,      41.695 , 41.541 , 41.545 , 41.593 , 41.692 , 42.239
RMSNormBenchmark ,Eager ,"RMSNormBenchmark_input_shape(2,128,1024)_musa_dtypetorch.float16_bwdall" ,"{\'input_shape\': [2, 128, 1024], \'device\': \'musa\', \'dtype\': \'torch.float16\'}" ,True     ,-1   ,     0 ,195.865 ,         322.567 ,     204.621 ,193.62  ,193.973 ,195.865 ,203.875 ,247.161
RMSNormBenchmark ,Eager ,"RMSNormBenchmark_input_shape(2,128,1024)_musa_dtypetorch.float16_bwd1"   ,"{\'input_shape\': [2, 128, 1024], \'device\': \'musa\', \'dtype\': \'torch.float16\'}" ,True     ,-1   ,     0 ,194.261 ,           0.13  ,     194.154 ,193.415 ,194.023 ,194.261 ,194.366 ,194.627
RMSNormBenchmark ,Eager ,"RMSNormBenchmark_input_shape(2,128,1024)_musa_dtypetorch.float16_bwd2"   ,"{\'input_shape\': [2, 128, 1024], \'device\': \'musa\', \'dtype\': \'torch.float16\'}" ,True     ,-1   ,     0 ,193.996 ,           0.188 ,     193.874 ,193.044 ,193.677 ,193.996 ,194.139 ,194.446
RMSNormBenchmark ,Eager ,"RMSNormBenchmark_input_shape(8,32,2048)_musa_dtypetorch.float16_bwdall"  ,"{\'input_shape\': [8, 32, 2048], \'device\': \'musa\', \'dtype\': \'torch.float16\'}"  ,True     ,-1   ,     0 ,214.838 ,           0.26  ,     214.932 ,214.025 ,214.703 ,214.838 ,215.33  ,215.59
RMSNormBenchmark ,Eager ,"RMSNormBenchmark_input_shape(8,32,2048)_musa_dtypetorch.float16_bwd1"    ,"{\'input_shape\': [8, 32, 2048], \'device\': \'musa\', \'dtype\': \'torch.float16\'}"  ,True     ,-1   ,     0 ,212.984 ,           0.376 ,     212.865 ,211.478 ,212.864 ,212.984 ,213.168 ,213.527
RMSNormBenchmark ,Eager ,"RMSNormBenchmark_input_shape(8,32,2048)_musa_dtypetorch.float16_bwd2"    ,"{\'input_shape\': [8, 32, 2048], \'device\': \'musa\', \'dtype\': \'torch.float16\'}"  ,True     ,-1   ,     0 ,212.607 ,           0.238 ,     212.679 ,211.908 ,212.349 ,212.607 ,213.086 ,213.368
RMSNormBenchmark ,Eager ,"RMSNormBenchmark_input_shape(8,128,2048)_musa_dtypetorch.float16_bwdall" ,"{\'input_shape\': [8, 128, 2048], \'device\': \'musa\', \'dtype\': \'torch.float16\'}" ,True     ,-1   ,     0 ,296.904 ,           0.083 ,     296.936 ,296.412 ,296.803 ,296.904 ,297.133 ,297.361
RMSNormBenchmark ,Eager ,"RMSNormBenchmark_input_shape(8,128,2048)_musa_dtypetorch.float16_bwd1"   ,"{\'input_shape\': [8, 128, 2048], \'device\': \'musa\', \'dtype\': \'torch.float16\'}" ,True     ,-1   ,     0 ,295.747 ,           0.275 ,     295.591 ,294.457 ,295.464 ,295.747 ,295.906 ,296.193
RMSNormBenchmark ,Eager ,"RMSNormBenchmark_input_shape(8,128,2048)_musa_dtypetorch.float16_bwd2"   ,"{\'input_shape\': [8, 128, 2048], \'device\': \'musa\', \'dtype\': \'torch.float16\'}" ,True     ,-1   ,     0 ,295.371 ,           0.024 ,     295.416 ,295.145 ,295.361 ,295.371 ,295.494 ,295.687
RMSNormBenchmark ,Eager ,"RMSNormBenchmark_input_shape(1,128,4096)_musa_dtypetorch.float16_bwdall" ,"{\'input_shape\': [1, 128, 4096], \'device\': \'musa\', \'dtype\': \'torch.float16\'}" ,True     ,-1   ,     0 ,221.199 ,           0.127 ,     221.237 ,220.813 ,220.886 ,221.199 ,221.591 ,221.692
```

--------------------------------

### Install PyTorch torchvision

Source: https://github.com/moorethreads/torch_musa/blob/main/README.md

Install torchvision from the official PyTorch repository. The version tag should correspond to your PyTorch version.

```shell
git clone https://github.com/pytorch/vision -b ${version} --depth 1
cd vision && python setup.py install
```

--------------------------------

### MUSA Kernel Implementation using CUDA Porting (Lerp Example)

Source: https://github.com/moorethreads/torch_musa/blob/main/docs/pdf/developer_guide/source/6_customize_operators/customize_structured_operators.md

Example of a MUSA kernel file generated by CUDA Porting for the lerp operator. It includes MUSA kernels for different groups and stub registrations.

```yaml
// 文件位置：build/generated_cuda_compatible/aten/src/ATen/native/musa/Lerp.mu

// lerp.Tensor group MUSA kernel
void lerp_tensor_kernel(at::TensorIteratorBase& iter) {......}

// lerp.Scalar group MUSA kernel
void lerp_scalar_kernel(at::TensorIteratorBase& iter, const c10::Scalar& weight) {
  ......
}

// lerp.Tensor group stub 注册
REGISTER_DISPATCH(lerp_kernel_tensor_weight, &lerp_tensor_kernel);
// lerp.Scalar group stub 注册
REGISTER_DISPATCH(lerp_kernel_scalar_weight, &lerp_scalar_kernel);
```

--------------------------------

### Using FusedAdamW Optimizer in torch_musa

Source: https://github.com/moorethreads/torch_musa/blob/main/torch_musa/optim/README.md

Instantiate and use FusedAdamW from torch_musa.optim for optimized performance on MUSA backend. This example shows the typical workflow of initializing the optimizer, zeroing gradients, performing a backward pass, and stepping the optimizer.

```python
from torch_musa.optim import FusedAdamW

optimizer = FusedAdamW(model.parameters(), lr=0.1, momentum=0.9)
optimizer.zero_grad()
loss_fn(model(input), target).backward()
optimizer.step()
```

--------------------------------

### Example Benchmark Output for RMSNorm

Source: https://github.com/moorethreads/torch_musa/blob/main/benchmark/operator_benchmark/README.md

This is an example of the output generated when running the RMSNorm benchmark test. It shows the execution time in microseconds for different input shapes, devices, and data types.

```plaintext
# ----------------------------------------
# PyTorch/Caffe2 Operator Micro-benchmarks
# ----------------------------------------
# Tag : short

# Benchmarking PyTorch: RMSNormBenchmark
# Mode: Eager
# Name: RMSNormBenchmark_input_shape(2,128,1024)_musa_dtypetorch.float16
# Input: input_shape: (2, 128, 1024), device: musa, dtype: torch.float16
Forward Execution Time (us) : 29.381

# Benchmarking PyTorch: RMSNormBenchmark
# Mode: Eager
# Name: RMSNormBenchmark_input_shape(2,128,1024)_musa_dtypetorch.float32
# Input: input_shape: (2, 128, 1024), device: musa, dtype: torch.float32
Forward Execution Time (us) : 30.456

# Benchmarking PyTorch: RMSNormBenchmark
# Mode: Eager
# Name: RMSNormBenchmark_input_shape(8,32,2048)_musa_dtypetorch.float16
# Input: input_shape: (8, 32, 2048), device: musa, dtype: torch.float16
Forward Execution Time (us) : 34.333

# Benchmarking PyTorch: RMSNormBenchmark
# Mode: Eager
# Name: RMSNormBenchmark_input_shape(8,32,2048)_musa_dtypetorch.float32
# Input: input_shape: (8, 32, 2048), device: musa, dtype: torch.float32
Forward Execution Time (us) : 35.437
```

--------------------------------

### Run Development Docker Container

Source: https://github.com/moorethreads/torch_musa/blob/main/docs/markdown/source/03_compile_and_install.md

Launches a Docker container for torch_musa development. Ensure mt-container-toolkit is installed and include '--env MTHREADS_VISIBLE_DEVICES=all' for MUSA device access.

```shell
docker run -it --name=torch_musa_dev --env MTHREADS_VISIBLE_DEVICES=all --shm-size=80g sh-harbor.mthreads.com/mt-ai/musa-pytorch-dev-py38:latest /bin/bash
```

--------------------------------

### Set Environment Variables for Building

Source: https://github.com/moorethreads/torch_musa/blob/main/README.md

Set essential environment variables for building torch_musa. Ensure MUSA_HOME points to your MUSA installation path.

```bash
export MUSA_HOME=path/to/musa/install/path  # defalut value is /usr/local/musa/
export LD_LIBRARY_PATH=$MUSA_HOME/lib:$LD_LIBRARY_PATH
export PATH=$MUSA_HOME/bin:$PATH

# if PYTORCH_REPO_PATH is not set, PyTorch will be downloaded outside this directory when building with build.sh
export PYTORCH_REPO_PATH=/path/to/PyTorch
```

--------------------------------

### Enable TensorCore with TF32

Source: https://github.com/moorethreads/torch_musa/blob/main/docs/markdown/source/04_quick_start.md

This example shows how to enable TensorFloat32 (TF32) for TensorCore acceleration on S4000 devices when using float32 inputs. TF32 can speed up matrix multiplication operations. The `torch.backends.mudnn.flags` context manager or a global setting can be used.

```python
import torch
import torch_musa
with torch.backends.mudnn.flags(allow_tf32=True):
    assert torch.backends.mudnn.allow_tf32
    a = torch.randn(10240, 10240, dtype=torch.float, device='musa')
    b = torch.randn(10240, 10240, dtype=torch.float, device='musa')
    result_tf32 = a @ b

torch.backends.mudnn.allow_tf32 = True
assert torch_musa._MUSAC._get_allow_tf32()
a = torch.randn(10240, 10240, dtype=torch.float, device='musa')
b = torch.randn(10240, 10240, dtype=torch.float, device='musa')
result_tf32 = a @ b
```

--------------------------------

### Enable File Logging in Distributed Settings

Source: https://github.com/moorethreads/torch_musa/blob/main/docs/markdown/source/08_debug.md

Use the `should_log_to_file` switch to prevent cross-rank interference in logs within a distributed setup. This ensures cleaner log output.

```python
from torch_musa.utils.compare_tool import open_module_tracker, ModuleInfo

model = get_your_model()
open_module_tracker(model)
with CompareWithCPU(atol=0.001, rtol=0.001, verbose=True, should_log_to_file=True, output_dir='path_to_save'):
    train(model)
```

--------------------------------

### Distributed Data Parallel (DDP) Training with Profiler

Source: https://github.com/moorethreads/torch_musa/blob/main/docs/pdf/developer_guide/source/9_optimize_performance/optimize_performance.md

This script demonstrates a distributed data parallel training setup using PyTorch and torch_musa, incorporating the profiler to capture CPU and MUSA activities. Ensure environment variables like MASTER_ADDR and MASTER_PORT are set for distributed training.

```python
import os
import argparse
import torch
from torch import nn
from torch import optim
from torch.nn.parallel import DistributedDataParallel as DDP
import torch.distributed as dist
import torch.multiprocessing as mp
import torch_musa
import torchvision
import torchvision.transforms as T
from torchvision import models

model = models.resnet50(weights=models.ResNet50_Weights.IMAGENET1K_V1)

transform = T.Compose([T.Resize(256), T.CenterCrop(224), T.ToTensor()])
trainset = torchvision.datasets.CIFAR10(root='./data', train=True,
                                        download=True, transform=transform)
trainloader = torch.utils.data.DataLoader(trainset, batch_size=32,
                                          shuffle=True, num_workers=4)
criterion = nn.CrossEntropyLoss()

def clean():
    dist.destroy_process_group()

def start(rank, world_size):
    if os.getenv("MASTER_ADDR") is None:
        os.environ['MASTER_ADDR'] = '127.0.0.1'
    if os.getenv("MASTER_PORT") is None:
        os.environ['MASTER_PORT'] = '29500'
    dist.init_process_group("mccl", rank=rank, world_size=world_size)

def runner(gpu, args):
    rank = args.nr * args.gpus + gpu
    torch_musa.set_device(rank % torch.musa.device_count())
    start(rank, args.world_size)
    model.to('musa')
    ddp_model = DDP(model, device_ids=[rank % torch.musa.device_count()])
    optimizer = optim.SGD(ddp_model.parameters(), lr=0.001)
    with torch.profiler.profile(
        activities=[
            torch.profiler.ProfilerActivity.CPU,
            torch.profiler.ProfilerActivity.MUSA],
        schedule=torch.profiler.schedule(
            wait=1,
            warmup=1,
            active=2),
        on_trace_ready=torch.profiler.tensorboard_trace_handler('./result_dist_ddp', worker_name='worker'+str(rank)),
        record_shapes=True,
        profile_memory=True,  # This will take 1 to 2 minutes. Setting it to False could greatly speedup.
        with_stack=True,
        experimental_config=torch._C._profiler._ExperimentalConfig(verbose=True)
    ) as p:
        for step, data in enumerate(trainloader, 0):
            inputs, labels = data[0].to('musa'), data[1].to('musa')
            outputs = ddp_model(inputs)
            loss = criterion(outputs, labels)
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
            p.step()
            if step + 1 >= 4:
                break
    clean()

def train(fn, args):
    mp.spawn(fn, args=(args,), nprocs=args.gpus, join=True)

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument('-n', '--nodes', default=1,
                        type=int, metavar='N')
    parser.add_argument('-g', '--gpus', default=1, type=int,
                        help='number of gpus per node')
    parser.add_argument('-nr', '--nr', default=0, type=int,
                        help='ranking within the nodes')
    args = parser.parse_args()
    args.world_size = args.gpus * args.nodes
    train(runner, args)

```

--------------------------------

### Run TorchMUSA within Docker

Source: https://github.com/moorethreads/torch_musa/blob/main/README.md

Example command to pull and run a TorchMUSA Docker image, setting up the environment for MUSA-accelerated PyTorch operations. Ensure MTHREADS_VISIBLE_DEVICES is set to 'all' for visibility.

```bash
# For example, start a S80 docker with Python3.10
docker run -it --privileged \
  --pull always --network=host \
  --env MTHREADS_VISIBLE_DEVICES=all \
  registry.mthreads.com/mcconline/musa-pytorch-release-public:latest /bin/bash
```

--------------------------------

### Basic PyTorch Inference Example

Source: https://github.com/moorethreads/torch_musa/blob/main/docs/markdown/source/04_quick_start.md

This snippet demonstrates a basic inference loop for a PyTorch model, calculating accuracy on test data. Ensure the model and data are moved to the correct device.

```python
import torch

# Assuming 'net', 'testloader', 'device', 'total', 'correct' are defined elsewhere
# Example placeholder definitions:
# class Net(torch.nn.Module): pass
# net = Net()
# device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# net.to(device)
# testloader = [(torch.randn(1, 3, 32, 32), torch.randint(0, 10, (1,)))] * 100
# total = 0
# correct = 0

# since we're not training, we don't need to calculate the gradients for our outputs
with torch.no_grad():
    for data in testloader:
        images, labels = data
        # calculate outputs by running images through the network
        outputs = net(images.to(device))
        # the class with the highest energy is what we choose as prediction
        _, predicted = torch.max(outputs.data, 1)
        total += labels.size(0)
        correct += (predicted == labels.to(device)).sum().item()
print(f'Accuracy of the network on the 10000 test images: {100 * correct // total} %')
```

--------------------------------

### AMP Training with torch_musa

Source: https://github.com/moorethreads/torch_musa/blob/main/docs/markdown/source/04_quick_start.md

Example of Automatic Mixed Precision (AMP) training using torch_musa. This enables faster training by using lower precision data types where appropriate. Ensure your model and data are on the 'musa' device.

```python
import torch
import torch_musa
import torch.nn as nn

class SimpleModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(5, 4)
        self.relu = nn.ReLU()
        self.fc2 = nn.Linear(4, 3)
    def forward(self, x):
        x = self.fc1(x)
        x = self.relu(x)
        x = self.fc2(x)
        return x
    def __call__(self, x):
        return self.forward(x)

DEVICE = "musa"

def train_in_amp(low_dtype=torch.float16):
    model = SimpleModel().to(DEVICE)
    criterion = nn.MSELoss()
    optimizer = torch.optim.SGD(model.parameters(), lr=0.1)

    # create the scaler object
    scaler = torch.musa.amp.GradScaler()

    inputs = torch.randn(6, 5).to(DEVICE) # 将数据移至 GPU
    targets = torch.randn(6, 3).to(DEVICE)
    for step in range(20):
        optimizer.zero_grad()
        # create autocast environment
        with torch.musa.amp.autocast(dtype=low_dtype):
            outputs = model(inputs
            assert outputs.dtype == low_dtype
            loss = criterion(outputs, targets)

        scaler.scale(loss).backward()
        scaler.step(optimizer)
        scaler.update()
    return loss
if __name__ == "__main__":
    train_in_amp(torch.float16)
```

--------------------------------

### Enable NHWC Layout Optimization for Convolutional Models

Source: https://github.com/moorethreads/torch_musa/blob/main/docs/pdf/developer_guide/source/9_optimize_performance/optimize_performance.md

This example demonstrates how to enable NHWC layout optimization for convolutional models on s4000 devices using torch_musa. This optimization can lead to performance improvements.

```python
import torch
import torch_musa

torch.backends.mudnn.allow_tf32 = True
model = Model() # define model here
model = model.to(memory_format=torch.channels_last) # transform layout to NHWC
```

--------------------------------

### Recommended Tensor Creation with Device Info

Source: https://github.com/moorethreads/torch_musa/blob/main/torch_musa/csrc/aten/ops/README.md

Prefer using `at::empty_like` with options that include device information to create tensors. This ensures that device index information is preserved.

```cpp
at::empty_like(input, input.options())
```

--------------------------------

### Build PDF Documentation

Source: https://github.com/moorethreads/torch_musa/blob/main/docs/pdf/README.md

Executes the build script to generate PDF documentation. This script assumes all dependencies are met.

```bash
bash build.sh
```

--------------------------------

### Install MUSA torchvision

Source: https://github.com/moorethreads/torch_musa/blob/main/README.md

Install torchvision from the MooreThreads repository for MUSA v2.7.0 and later. Ensure you are using the correct branch for your MUSA version.

```shell
git clone https://github.com/MooreThreads/vision -b v0.24.1-musa --depth 1
cd vision && python setup.py install
```

--------------------------------

### Build Torch-MUSA from Source

Source: https://github.com/moorethreads/torch_musa/blob/main/README.md

Build PyTorch and torch_musa from scratch using the build script. Use the -c flag to clean the cache before building.

```bash
bash build.sh -c
```

--------------------------------

### MMCV MUSA Extension Build Logic

Source: https://github.com/moorethreads/torch_musa/blob/main/docs/markdown/source/06_third_party_lib_extension.md

This snippet demonstrates the adaptation of MMCV's setup.py for MUSA compatibility, including environment variable checks, MUSA-specific imports, and path configurations.

```python
elif os.getenv('FORCE_MUSA', '0') == '1':
  from torch_musa.utils.simple_porting import SimplePorting
  from torch_musa.utils.musa_extension import MUSAExtension
  SimplePorting(cuda_dir_path="./mmcv/ops/csrc", ignore_dir_paths=[
  "./mmcv/ops/csrc/common/mlu",
  "./mmcv/ops/csrc/common/mps",
  "./mmcv/ops/csrc/parrots",
  "./mmcv/ops/csrc/pytorch/mlu",
  "./mmcv/ops/csrc/pytorch/mps",
  "./mmcv/ops/csrc/pytorch/npu"
  ],
  mapping_rule={
      "_CU_H_": "_MU_H_",
      "_CUH": "_MUH",
      "__NVCC__": "__MUSACC__",
      "MMCV_WITH_CUDA": "MMCV_WITH_MUSA",
      "AT_DISPATCH_FLOATING_TYPES_AND_HALF": "AT_DISPATCH_FLOATING_TYPES",
      "#include <ATen/cuda/CUDAContext.h>": "#include \"torch_musa/csrc/aten/musa/MUSAContext.h\"",
      "#include <c10/cuda/CUDAGuard.h>": "#include \"torch_musa/csrc/core/MUSAGuard.h\"",
      "::cuda::": "::musa::",
      "/cuda/": "/musa/",
      ", CUDA,": ", PrivateUse1,",
      ".cuh": ".muh",
      ".is_cuda()": ".is_privateuseone()",
      }
      ).run()
  op_files = glob.glob('./mmcv/ops/csrc_musa/pytorch/*.cpp') + \
      glob.glob('./mmcv/ops/csrc_musa/pytorch/cpu/*.cpp') + \
      glob.glob('./mmcv/ops/csrc_musa/pytorch/cuda/*.mu') + \
      glob.glob('./mmcv/ops/csrc_musa/pytorch/cuda/*.cpp')
  from torch_musa.testing import get_musa_arch
  define_macros += [('MMCV_WITH_MUSA', None),
                    ('MUSA_ARCH', str(get_musa_arch()))]
  os.environ['MUSA_ARCH'] = str(get_musa_arch())

  extension = MUSAExtension
  include_dirs.append(os.path.abspath('./mmcv/ops/csrc_musa/pytorch'))
  include_dirs.append(os.path.abspath('./mmcv/ops/csrc_musa/common'))
  include_dirs.append(os.path.abspath('./mmcv/ops/csrc_musa/common/cuda'))
  from torch_musa.utils.musa_extension import MUSAExtension,BuildExtension
  cmd_class = {'build_ext': BuildExtension}
elif (hasattr(torch, 'is_mlu_available') and
...
```

--------------------------------

### Finding PyTorch and TorchMusa Packages

Source: https://github.com/moorethreads/torch_musa/blob/main/examples/cpp/kernel/CMakeLists.txt

Locates the installed PyTorch and TorchMusa libraries. This step is critical for linking against PyTorch functionalities.

```cmake
find_package(Torch REQUIRED)
if(NOT Torch_FOUND)
  message(
    FATAL_ERROR
      "PyTorch/LibTorch not found. Ensure it's installed and CMAKE_PREFIX_PATH is set correctly."
  )
endif()
message(STATUS "TORCH_LIBRARIES: ${TORCH_LIBRARIES}")

find_package(TorchMusa REQUIRED CONFIG)
```

--------------------------------

### Get Benchmark Results at Runtime

Source: https://github.com/moorethreads/torch_musa/blob/main/benchmark/operator_benchmark/README.md

Retrieve benchmark results programmatically at runtime using the BenchmarkRunner API. The results are returned as a dictionary.

```python
# create a benchmark runner 
runner = BenchmarkRunner(*args)
benchmark_result = runner.run()
```

--------------------------------

### Distributed Training with torch_musa

Source: https://github.com/moorethreads/torch_musa/blob/main/docs/markdown/source/04_quick_start.md

Demonstrates distributed training using PyTorch's DistributedDataParallel (DDP) with torch_musa. This requires setting environment variables for master address and port, and initializing the process group. The `mp.spawn` function is used to launch multiple processes.

```python
"""Demo of DistributedDataParall"""
import os
import torch
from torch import nn
from torch import optim
from torch.nn.parallel import DistributedDataParallel as DDP
import torch.distributed as dist
import torch.multiprocessing as mp
import torch_musa

class Model(nn.Module):
    def __init__(self):
        super().__init__()
        self.linear = nn.Linear(5,5)
    def forward(self, x):
        return self.linear(x)

def start(rank, world_size):
    if os.getenv("MASTER_ADDR") is None:
        os.environ["MASTER_ADDR"]= ip # IP must be specified here
    if os.getenv("MASTER_PORT") is None:
        os.environ["MASTER_PORT"]= port # port must be specified here
    dist.init_process_group("mccl", rank=rank, world_size=world_size)

def clean():
    dist.destroy_process_group()    

def runner(rank, world_size):
    torch_musa.set_device(rank)
    start(rank, world_size)
    model = Model().to('musa')
    ddp_model = DDP(model, device_ids=[rank])
    optimizer = optim.SGD(ddp_model.parameters(), lr=0.001)
    for _ in range(5):
        input_tensor = torch.randn(5, dtype=torch.float, requires_grad=True).to('musa')
        target_tensor = torch.zeros(5, dtype=torch.float).to('musa')
        output_tensor = ddp_model(input_tensor)
        loss_f = nn.MSELoss()
        loss = loss_f(output_tensor, target_tensor)
        loss.backward()
        optimizer.step()
    clean()

if __name__ == "__main__":
    mp.spawn(runner, args=(2,), nprocs=2, join=True)
```

--------------------------------

### Create Render Group and Add User

Source: https://github.com/moorethreads/torch_musa/blob/main/docs/markdown/source/10_FAQ.md

If the 'render' group does not exist, create it with a specific GID and then add the current user to this group. This is a prerequisite for non-root users to access graphics card information.

```shell
sudo groupadd -o -g 109 render
sudo usermod -aG render `whoami`
```

--------------------------------

### Build and Link TorchMusa Application

Source: https://github.com/moorethreads/torch_musa/blob/main/examples/cpp/CMakeLists.txt

Defines the executable target 'example-app' and links it against the torch_musa library and libtinfo.so.6. Ensure that the paths to libraries are correct for your system.

```cmake
add_executable(example-app example-app.cpp)
target_link_libraries(example-app torch_musa
                      "/usr/lib/x86_64-linux-gnu/libtinfo.so.6")
```

--------------------------------

### Functional Rule Invocation (tril)

Source: https://github.com/moorethreads/torch_musa/blob/main/docs/pdf/developer_guide/source/6_customize_operators/operator_categories.md

Example of invoking the tril operator using the default functional rule, where the result tensor is created internally.

```python
>>> a = torch.randn(3, 3)
>>> torch.tril(a)
```

--------------------------------

### Build Torch-MUSA and Generate Wheels

Source: https://github.com/moorethreads/torch_musa/blob/main/README.md

Build only torch_musa and generate wheels, assuming PyTorch is already built. Use the -m and -w flags.

```bash
bash build.sh -m -w
```

--------------------------------

### Example of tril Operator Execution

Source: https://github.com/moorethreads/torch_musa/blob/main/docs/pdf/developer_guide/source/6_customize_operators/when_need_to_customize_operators.md

This Python code demonstrates how to execute the tril operator. If torch_musa has implemented the operator, the program will complete successfully.

```python
import torch
import torch_musa
input_data = torch.randn(3, 3, device="musa")
result = torch.tril(input_data)
```

--------------------------------

### Customized Strategy: Full Kernel Implementation Customization

Source: https://github.com/moorethreads/torch_musa/blob/main/tools/codegen/README.md

This strategy enables full customization of the kernel implementation, including both preprocessing (`meta`) and calculation (`impl`) logic. Both `meta` and `impl` functions are generated under the `at::musa` namespace.

```yaml
- func: bitwise_right_shift.Tensor      # fucntional overload
- func: bitwise_right_shift_.Tensor     # inplace overload
- func: bitwise_right_shift.Tensor_out  # out overload
  structured_inherits: MyBase           # optional
  precomputed:                          # optional
  - dim -> int dim1, int dim2
  dispatch:
    PrivateUse1: musa_bitwise_right_shift_impl
```

--------------------------------

### Set Environment Variables for Compilation

Source: https://github.com/moorethreads/torch_musa/blob/main/docs/markdown/source/03_compile_and_install.md

Configures necessary environment variables for the build process, including MUSA installation path and PyTorch source directory.

```shell
export MUSA_HOME=path/to/musa_libraries(including musa_toolkits, mudnn and so on) # defalut value is /usr/local/musa/ 
export LD_LIBRARY_PATH=$MUSA_HOME/lib:$LD_LIBRARY_PATH
export PYTORCH_REPO_PATH=path/to/PyTorch source code
# if PYTORCH_REPO_PATH is not set, PyTorch-v2.0.0 will be downloaded outside this directory automatically when building with build.sh
```

--------------------------------

### C++ Deployment with Torch-MUSA

Source: https://github.com/moorethreads/torch_musa/blob/main/docs/pdf/developer_guide/source/5_quickguide/quickguide.md

This C++ code snippet illustrates how to load and run a Torch JIT traced or scripted model on a MUSA device. It includes registering the 'musa' backend and preparing input tensors.

```cpp
#include <torch/script.h>
#include <torch_musa/csrc/core/Device.h>
#include <iostream>
#include <memory>

int main(int argc, const char* argv[]) {
  // Register 'musa' for PrivateUse1 as we save model with 'musa'.
  c10::register_privateuse1_backend("musa");

  torch::jit::script::Module module;
  // Load model which saved with torch jit.trace or jit.script.
  module = torch::jit::load(argv[1]);

  std::vector<torch::jit::IValue> inputs;
  // Ready for input data.
  torch::Tensor input = torch::rand({1, 3, 224, 224}).to("musa");
  inputs.push_back(input);

  // Model execute.
  at::Tensor output = module.forward(inputs).toTensor();

  return 0;
}
```

--------------------------------

### Enable File Logging in Distributed Mode

Source: https://github.com/moorethreads/torch_musa/blob/main/docs/pdf/developer_guide/source/10_debug_tools/debug_tools.md

Use the `should_log_to_file` flag to prevent cross-rank interference in logs when using `CompareWithCPU` or `NanInfTracker` in a distributed setup.

```python
from torch_musa.utils.compare_tool import open_module_tracker, ModuleInfo

model = get_your_model()
open_module_tracker(model)
with CompareWithCPU(atol=0.001, rtol=0.001, verbose=True, should_log_to_file=True, output_dir='path_to_save'):
    train(model)

```

--------------------------------

### Format CMake Files with cmakelang

Source: https://github.com/moorethreads/torch_musa/blob/main/tools/formatting/README.md

Use this command to format CMake files in the torch_musa project. Ensure you have installed the necessary requirements and have the .cmake-format.yaml configuration file present.

```bash
cd torch_musa
pip install -r requirements.txt
cmake-format -c ./.cmake-format.yaml -i ./CMakeLists.txt  ./torch_musa/csrc/CMakeLists.txt
```

```bash
pre-commit run cmake-format --all-files
```

--------------------------------

### Set MUSA_HOME Environment Variable

Source: https://github.com/moorethreads/torch_musa/blob/main/examples/cpp/CMakeLists.txt

Sets the MUSA_HOME environment variable to specify the installation path of MUSA. This is often required for build systems to locate MUSA components.

```cmake
set(ENV{MUSA_HOME} "/usr/local/musa")
```

--------------------------------

### Dispatch Unstructured Kernels to MUSA

Source: https://github.com/moorethreads/torch_musa/blob/main/tools/codegen/README.md

For PyTorch unstructured kernels, specify the implementation function in the `dispatch` key with the MUSA backend. This example shows how to dispatch a convolution overrideable kernel.

```yaml
- func: convolution_overrideable
  dispatch:
    PrivateUse1: Convolution
```

--------------------------------

### Defining Build Targets and Linking Libraries

Source: https://github.com/moorethreads/torch_musa/blob/main/examples/cpp/kernel/CMakeLists.txt

Defines the executable target 'qlinear' and links it with necessary libraries including PyTorch, MUSA runtime, and TorchMusa. This step compiles the source files into a runnable program.

```cmake
set(MU_SRCS MatmulQuantized.mu)
set(CPP_SRCS Qlinear.cpp)
set(CMAKE_BUILD_RPATH
    "/usr/local/lib/python3.10/dist-packages/torch_musa-2.7.1-py3.10-linux-x86_64.egg/torch_musa/lib"
)
musa_add_executable(qlinear ${MU_SRCS} ${CPP_SRCS})
target_link_libraries(
  qlinear
  torch_musa
  musart
  ${TORCH_LIBRARIES}
  torch_cpu)
```

--------------------------------

### Add User to Render Group

Source: https://github.com/moorethreads/torch_musa/blob/main/docs/markdown/source/10_FAQ.md

If a non-root user cannot view the graphics card after installing the driver, add the user to the 'render' group. This command adds the current user to the 'render' group.

```shell
sudo usermod -aG render `whoami`
```

--------------------------------

### Advanced SimplePorting CLI options

Source: https://github.com/moorethreads/torch_musa/blob/main/torch_musa/utils/README.md

Customize SimplePorting with options like ignoring directories, custom mapping rules, dropping default mappings, and specifying a mapping directory.

```bash
python -m torch_musa.utils.simple_porting --cuda-dir-path cuda/ --ignore-dir-paths ["csrc/npu"] --mapping-rule {"cuda":"musa"} --drop-default-mapping --mapping-dir-path mapping/
```

--------------------------------

### Legacy Strategy: Basic Kernel Registration

Source: https://github.com/moorethreads/torch_musa/blob/main/tools/codegen/README.md

Use this strategy for the simplest kernel registration where MUSA kernels are implemented by porting CUDA and registered into `{op}_stub`. The actual dispatch happens inside the `impl` function, with preprocessing and calculation logic identical to CPU/CUDA.

```yaml
- func: bitwise_right_shift.Tensor      # fucntional overload
- func: bitwise_right_shift_.Tensor     # inplace overload
- func: bitwise_right_shift.Tensor_out  # out overload
```

--------------------------------

### Example of Operator Implementation in CUDA (abs.out)

Source: https://github.com/moorethreads/torch_musa/blob/main/docs/pdf/developer_guide/source/6_customize_operators/support_new_operators.md

This snippet shows a typical implementation of an operator in CUDA, which might be adapted for MUSA using CUDA-Porting. It demonstrates how to wrap a native ATen function.

```c++
at::Tensor & wrapper_CUDA_out_abs_out(const at::Tensor & self, at::Tensor & out) {
  // No device check
const OptionalDeviceGuard device_guard(device_of(self));
return at::native::abs_out(self, out);
}

******
m.impl("abs.out", TORCH_FN(wrapper_CUDA_out_abs_out));
```

--------------------------------

### Control Comparison Steps in AMP Scenarios

Source: https://github.com/moorethreads/torch_musa/blob/main/torch_musa/utils/README.md

Manage the activation of comparisons or NaN/Inf tracking in AMP scenarios using `start_step` and `end_step`. Comparisons are active only when `start_step <= step_cnt < end_step`.

```python
from torch_musa.utils.compare_tool import open_module_tracker, ModuleInfo

model = get_your_model()
open_module_tracker(model)
with CompareWithCPU(atol=0.001, rtol=0.001, verbose=True, start_step=5) as compare_with_cpu:
    for epoch in range(epoch_num):
        for step in range(step_num):
             train_step(model)
             compare_with_cpu.step()
```

--------------------------------

### Functional Rule Example (isnan)

Source: https://github.com/moorethreads/torch_musa/blob/main/docs/pdf/developer_guide/source/6_customize_operators/operator_categories.md

Demonstrates the functional calling rule where the operator creates and returns a new tensor for the result. This is common when type conversions or internal temporary variables are needed.

```python
torch.isnan(input) → Tensor
```