Add documentation for Huawei NPU supports

Author: ljq12697

1

@@ -10,7 +10,7 @@
 
 DNALLM-Suite is a comprehensive, open-source toolkit designed for fine-tuning and inference with DNA Language Models. It provides a unified interface for working with various DNA sequence models, supporting tasks ranging from basic sequence classification to advanced in-silico mutagenesis analysis. With built-in Model Context Protocol (MCP) support, DNALLM-Suite enables seamless communication with traditional large language models, allowing for enhanced integration and interoperability in AI-powered DNA analysis workflows.
 
-## 📦 Installation
+## 📦 Quick Installation
 
 ```bash
 pip install dnallm
@@ -48,6 +48,23 @@ DNALLM-Suite supports a wide range of DNA language models including:
 - **ModelScope**: Alternative model source with additional models
 - **Custom Models**: Support for locally trained or custom architectures
 
+## 🖥️ Supported Platforms
+
+DNALLM-Suite has been tested on a wide range of platforms and devices:
+
+### Platforms:
+- **Linux**
+- **Windows**
+- **MacOS**
+
+### Device:
+- **CPU**
+- **Nvidia GPU (CUDA)**
+- **AMD GPU (ROCm)**
+- **Apple Silicon (MPS)**
+- **Huawei Ascend NPU (CANN)**
+- **Intel Arc GPU (XPU)**
+
 ## 🛠️ Installation
 
 ### Prerequisites

README.md

@@ -150,7 +150,7 @@ For development and testing without GPU acceleration:
 
 ```bash
 # Create environment
-conda create -n dnallm-cpu python=3.12 -y
+conda create -n dnallm-cpu python=3.12 uv -y
 conda activate dnallm-cpu
 
 # Install all dependencies and CPU version
@@ -169,7 +169,7 @@ For GPU-accelerated training and inference:
 nvidia-smi
 
 # Create environment (using CUDA 12.4 as example)
-conda create -n dnallm-gpu python=3.12 -y
+conda create -n dnallm-gpu python=3.12 uv -y
 conda activate dnallm-gpu
 
 # Install all dependencies and CUDA 12.4 support
@@ -179,7 +179,61 @@ uv pip install -e '.[all,cuda124]'
 python -c "import torch; print(f'PyTorch: {torch.__version__}'); print(f'CUDA available: {torch.cuda.is_available()}')"
 ```
 
-### Scenario 3: Using Mamba Model Architecture
+### Scenario 3: Using Huawei Ascend NPU for Training and Inference
+
+For Huawei Ascend NPU accelerated training and inference, users should first check their device and environment, then install the appropriate dependencies.
+
+If Huawei Ascend driver is not installed in the machine, first check the device and install the corresponding drivers.
+
+For NPU driver, please refer to: https://www.hiascend.com/hardware/firmware-drivers/community
+
+For Ascend Extension for PyTorch (CANN driver), please refer to: https://www.hiascend.com/zh/cann/download
+
+For example, if you have a `Ascend 910B NPU with AArch64 architecture`, install drivers like this:
+
+```bash
+# Install NPU driver
+wget -c "https://ascend-repo.obs.cn-east-2.myhuaweicloud.com/Ascend%20HDK/Ascend%20HDK%2025.5.2/Ascend-hdk-910b-npu-driver_25.5.2_linux-aarch64.run"
+bash ./Ascend-hdk-910b-npu-driver_25.5.2_linux-aarch64.run --install
+# Install CANN Toolkit
+wget https://ascend-cann-open.obs.cn-north-4.myhuaweicloud.com/CANN/CANN%209.0.0/Ascend-cann_9.0.0_linux-aarch64.run
+bash ./Ascend-cann_9.0.0_linux-aarch64.run --install
+# Install CANN Kernel/Ops
+wget https://ascend-repo.obs.cn-east-2.myhuaweicloud.com/CANN/CANN%209.0.0/Ascend-cann-910b-ops_9.0.0_linux-aarch64.run
+bash ./Ascend-cann-910b-ops_9.0.0_linux-aarch64.run --install
+
+# Check the drivers
+source /usr/local/Ascend/cann/set_env.sh
+python3 -c "import acl;print(acl.get_soc_name())"
+npu-smi info
+```
+
+if you want to auto-activate the driver environment, add the `set_env.sh` to system environment.
+```bash
+echo "source /usr/local/Ascend/cann/set_env.sh" >> ~/.bashrc
+```
+
+To use the NPU accelerating in torch, a specific version of `torch_npu` package is also required. Please refer to [this page](https://gitcode.com/Ascend/pytorch) to check the dependency map.
+
+For example, CANN 9.0.0 support Pytorch version from 2.7.1 to 2.10.0, also the Python version need to >=3.9.
+
+```bash
+# Create environment (using CANN 9.0.0 as example)
+conda create -n dnallm-npu python=3.11 uv -y
+conda activate dnallm-npu
+
+# Install dependencies of NPU support
+uv pip install torch torch_npu==2.9.0
+
+# Verify installation
+python -c "import torch; import torch_npu; print(f'PyTorch: {torch.__version__}'); print(f'NPU available: {torch.npu.is_available()}')"
+```
+
+During training or inference, Huawei Ascend NPU accelerate is supported for most of the DNA models (models supported by Huggingface Transformers library).
+
+For other non-transformer models or CUDA-dependent models, Huawei provides a specific framework for efficient model training and inference, named [MindSpeed](https://gitcode.com/Ascend/MindSpeed-LLM/). Detailed supported model list is shown [here](https://gitcode.com/Ascend/MindSpeed-LLM/blob/master/docs/zh/pytorch/models/supported_models.md).
+
+### Scenario 4: Using Mamba Model Architecture
 
 For models with Mamba architecture (Plant DNAMamba, Caduceus, Jamba-DNA):
 
@@ -198,7 +252,7 @@ uv pip install -e '.[cuda124,mamba]' --no-cache-dir --no-build-isolation
 python -c "from mambapy import Mamba; print('Mamba installed successfully!')"
 ```
 
-### Scenario 4: Complete Development Environment
+### Scenario 5: Complete Development Environment
 
 For contributors and developers:
 
@@ -220,7 +274,7 @@ print('CUDA:', torch.version.cuda if torch.cuda.is_available() else 'CPU')
 "
 ```
 
-### Scenario 5: Running MCP Server Only
+### Scenario 6: Running MCP Server Only
 
 For MCP server deployment:
 

docs/getting_started/installation.md

@@ -99,6 +99,23 @@ uv pip install -e '.[cuda124]'
 uv pip install -e '.[cuda121]'
 ```
 
+### Huawei NPU Support
+
+For Huawei Ascend NPU acceleration, install the appropriate torch and extension version:
+
+```bash
+# For venv users: activate virtual environment
+source .venv/bin/activate
+# For conda users: activate conda environment
+# conda activate dnallm
+
+# Install specific torch and extension version according the CANN driver's version
+uv pip install torch torch_npu
+```
+
+Detailed installation method see [installation.md](./installation.md##Installation-Scenarios)
+
+
 ### Native Mamba Support
 
 Native Mamba architecture runs significantly faster than transformer-compatible Mamba architecture, but native Mamba depends on Nvidia GPUs.

docs/getting_started/quick_start.md

@@ -36,12 +36,26 @@ DNALLM-Suite supports a wide range of DNA language models including:
 - **ModelScope**: Alternative model source with additional models
 - **Custom Models**: Support for locally trained or custom architectures
 
+## 🖥️ Supported Platforms
+
+DNALLM-Suite has been tested on a wide range of platforms and devices:
+
+### Platforms
+- [**Linux**]  [**Windows**]  [**MacOS**]
+
+### Device
+- **CPU**
+- **Nvidia GPU (CUDA)**
+- **AMD GPU (ROCm)**
+- **Apple Silicon (MPS)**
+- **Huawei Ascend NPU (CANN)**
+- **Intel Arc GPU (XPU)**
+
 ## 🚀 Quick Start
 
 DNALLM-Suite uses conda for environment management and uv for dependency management and packaging.
 
 1. **Install dependencies (recommended: [uv](https://docs.astral.sh/uv/))**
-
 ```bash
 # Clone repository
 git clone https://github.com/zhangtaolab/DNALLM.git
@@ -69,6 +83,9 @@ uv pip install -e '.[all]'
 python -c "import dnallm; print('DNALLM installed successfully!')"
 ```
 
+For NPU support such as Huawei Ascend, see [installation](getting_started/installation.md#scenario-3-using-huawei-ascend-npu-for-training-and-inference) section.
+
+
 2. **Basic Model Loading and Inference**
 ```python
 from dnallm import load_config, load_model_and_tokenizer, DNAInference
@@ -107,6 +124,8 @@ predictions = mutagenesis.evaluate(strategy="mean")
 # Visualize results
 plot = mutagenesis.plot(predictions, save_path="mutation_effects.pdf")
 ```
+
+4. **Model Fine-tuning**
 ```python
 from dnallm.datahandling import DNADataset
 from dnallm.finetune import DNATrainer
@@ -125,7 +144,8 @@ trainer = DNATrainer(model=model, config=configs, datasets=dataset)
 # Start training
 trainer.train()
 ```
-<!-- skip-verify: requires async event loop context -->
+
+5. **MCP Server Deployment**
 ```python
 # Start MCP server for real-time DNA sequence prediction
 from dnallm.mcp import DNALLMMCPServer
@@ -138,55 +158,55 @@ await server.initialize()
 server.start_server(host="0.0.0.0", port=8000, transport="streamable-http")
 ```
 
-6. **Launch Jupyter Lab, Marimo or Gradio App for interactive development:**
-
-#### Interactive Demos (Marimo)
-```bash
-# Fine-tuning demo
-uv run --no-sync marimo run example/marimo/finetune/finetune_demo.py
-
-# Inference demo
-uv run --no-sync marimo run example/marimo/inference/inference_demo.py
-
-# Benchmark demo
-uv run --no-sync marimo run example/marimo/benchmark/benchmark_demo.py
-```
-
-#### Jupyter Notebooks
-```bash
-# Launch Jupyter Lab
-uv run --no-sync jupyter lab
-
-# Available notebooks:
-# - example/notebooks/finetune_binary/ - Binary classification fine-tuning
-# - example/notebooks/finetune_multi_labels/ - Multi-label classification
-# - example/notebooks/finetune_NER_task/ - Named Entity Recognition
-# - example/notebooks/inference/ - Model inference
-# - example/notebooks/in_silico_mutagenesis/ - Mutation analysis
-# - example/notebooks/inference_for_tRNA/ - tRNA-specific analysis
-# - example/notebooks/generation_evo_models/ - EVO model inference
-# - example/notebooks/lora_finetune_inference/ - LoRA fine-tuning
-# - example/notebooks/embedding_attention.ipynb - Embedding and attention analysis
-# - example/notebooks/finetune_custom_head/ - Custom classification head
-# - example/notebooks/finetune_generation/ - Sequence generation
-# - example/notebooks/generation/ - Sequence generation examples
-# - example/notebooks/generation_megaDNA/ - MegaDNA model inference
-# - example/notebooks/interpretation/ - Model interpretation
-# - example/notebooks/data_prepare/ - Data preparation examples
-# - example/notebooks/benchmark/ - Model evaluation and benchmarking
-```
-
-#### Web-based UI (Gradio)
-```bash
-# Launch Gradio configuration generator app
-uv run --no-sync python ui/run_config_app.py
-
-# Or run the model config generator directly
-uv run --no-sync python ui/model_config_generator_app.py
-
-# For Generation, we also provide a app
-uv run --no-sync python ui/generation_task_app.py
-```
+6. **Launch Jupyter Lab, Marimo or Gradio App for interactive development**
+
+    (1) Interactive Demos (Marimo)
+    ```bash
+    # Fine-tuning demo
+    uv run --no-sync marimo run example/marimo/finetune/finetune_demo.py
+
+    # Inference demo
+    uv run --no-sync marimo run example/marimo/inference/inference_demo.py
+
+    # Benchmark demo
+    uv run --no-sync marimo run example/marimo/benchmark/benchmark_demo.py
+    ```
+
+    (2) Jupyter Notebooks
+    ```bash
+    # Launch Jupyter Lab
+    uv run --no-sync jupyter lab
+
+    # Available notebooks:
+    # - example/notebooks/finetune_binary/ - Binary classification fine-tuning
+    # - example/notebooks/finetune_multi_labels/ - Multi-label classification
+    # - example/notebooks/finetune_NER_task/ - Named Entity Recognition
+    # - example/notebooks/inference/ - Model inference
+    # - example/notebooks/in_silico_mutagenesis/ - Mutation analysis
+    # - example/notebooks/inference_for_tRNA/ - tRNA-specific analysis
+    # - example/notebooks/generation_evo_models/ - EVO model inference
+    # - example/notebooks/lora_finetune_inference/ - LoRA fine-tuning
+    # - example/notebooks/embedding_attention.ipynb - Embedding and attention analysis
+    # - example/notebooks/finetune_custom_head/ - Custom classification head
+    # - example/notebooks/finetune_generation/ - Sequence generation
+    # - example/notebooks/generation/ - Sequence generation examples
+    # - example/notebooks/generation_megaDNA/ - MegaDNA model inference
+    # - example/notebooks/interpretation/ - Model interpretation
+    # - example/notebooks/data_prepare/ - Data preparation examples
+    # - example/notebooks/benchmark/ - Model evaluation and benchmarking
+    ```
+
+    (3) Web-based UI (Gradio)
+    ```bash
+    # Launch Gradio configuration generator app
+    uv run --no-sync python ui/run_config_app.py
+
+    # Or run the model config generator directly
+    uv run --no-sync python ui/model_config_generator_app.py
+
+    # For Generation, we also provide a app
+    uv run --no-sync python ui/generation_task_app.py
+    ```
 
 ## 🎯 Supported Task Types