Starrick Liu
421db17d53
This release introduces significant performance optimizations, memory efficiency improvements, and enhanced system robustness: 🚀 Performance Optimizations: - Add three new fused CUDA kernels (rope_index, rot_pos_emb, get_window_index) for accelerated multimodal preprocessing - Implement FSDP2 support for distributed training with improved memory efficiency - Add Torch.compile integration for additional performance gains - Optimize memory usage: reduce peak allocation from 48GB to 24GB on 8-GPU setup 🔧 System Robustness: - Fix missing token position inputs in prediction pipeline - Add type-robust negation operations in RoPE CUDA kernels (half/bfloat16 support) - Fix dataset root parameter initialization in LeRobot data loader - Enhanced error handling and input validation across fusion operators 📚 Documentation & Usability: - Add comprehensive memory usage benchmarks and hardware recommendations - Update citation format with proper arXiv reference - Improve training configuration documentation with quick start guide - Add detailed API documentation for new fusion operators 🛠️ Technical Details: - Version bump to 1.0.1 - New CUDA kernels: rope_index.cu, rot_pos.cu, window_index.cu - FSDP2 state dict loading with distribute_tensor support - Enhanced multimodal RoPE with 3D position encoding - Window attention optimization for Vision Transformers Breaking Changes: None - all changes are backward compatible
64 lines
1.7 KiB
Python
64 lines
1.7 KiB
Python
import os
|
|
import torch
|
|
from pathlib import Path
|
|
from setuptools import setup, find_packages
|
|
from torch.utils.cpp_extension import BuildExtension, CUDAExtension
|
|
|
|
cwd = Path(os.path.dirname(os.path.abspath(__file__)))
|
|
|
|
nvcc_flags = [
|
|
"-std=c++17", # NOTE: CUTLASS requires c++17
|
|
"-DENABLE_BF16", # Enable BF16 for cuda_version >= 11
|
|
]
|
|
|
|
env_arch_list = os.environ.get("TORCH_CUDA_ARCH_LIST", None)
|
|
|
|
if env_arch_list:
|
|
# Let PyTorch builder to choose device to target for.
|
|
device_capability = ""
|
|
else:
|
|
device_capability = torch.cuda.get_device_capability()
|
|
device_capability = f"{device_capability[0]}{device_capability[1]}"
|
|
|
|
if device_capability:
|
|
nvcc_flags.extend(
|
|
[
|
|
f"--generate-code=arch=compute_{device_capability},code=sm_{device_capability}",
|
|
f"-DGROUPED_GEMM_DEVICE_CAPABILITY={device_capability}",
|
|
]
|
|
)
|
|
|
|
ext_modules = [
|
|
CUDAExtension(
|
|
"wallx_csrc",
|
|
[
|
|
"csrc/ops.cu",
|
|
"csrc/dual_asym_grouped_gemm.cu",
|
|
"csrc/permute.cu",
|
|
"csrc/rope.cu",
|
|
"csrc/rope_index.cu",
|
|
"csrc/rot_pos.cu",
|
|
"csrc/window_index.cu",
|
|
],
|
|
include_dirs=[f"{cwd}/3rdparty/cutlass/include/", f"{cwd}/csrc"],
|
|
extra_compile_args={
|
|
"cxx": ["-fopenmp", "-fPIC", "-Wno-strict-aliasing"],
|
|
"nvcc": nvcc_flags,
|
|
},
|
|
)
|
|
]
|
|
|
|
setup(
|
|
name="wall_x",
|
|
version="1.0.1",
|
|
author="X2Robot Team",
|
|
classifiers=[
|
|
"Programming Language :: Python :: 3",
|
|
"License :: OSI Approved :: BSD License",
|
|
"Operating System :: Unix",
|
|
],
|
|
packages=find_packages(),
|
|
ext_modules=ext_modules,
|
|
cmdclass={"build_ext": BuildExtension},
|
|
)
|