mirror of
https://github.com/modelscope/modelscope.git
synced 2026-07-12 21:39:26 +02:00
51a1b76e91
Link: https://code.alibaba-inc.com/Ali-MaaS/MaaS-lib/codereview/14678226 * modify librosa version * fix python3.10 compatible issue * remove healpy in requirements for windowns compatible
30 lines
733 B
Python
30 lines
733 B
Python
import os
|
|
|
|
import torch
|
|
|
|
from modelscope.utils.megatron_utils import init_megatron_util
|
|
|
|
|
|
def pre_compile_megatron_util():
|
|
dummy_megatron_cfg = {
|
|
'tensor_model_parallel_size': 1,
|
|
'world_size': 1,
|
|
'distributed_backend': 'nccl',
|
|
'seed': 42,
|
|
}
|
|
os.environ['MASTER_PORT'] = '39501'
|
|
init_megatron_util(dummy_megatron_cfg)
|
|
|
|
|
|
def pre_compile_all():
|
|
if torch.cuda.is_available(): # extension require cuda.
|
|
# pre compile pai-easycv
|
|
from easycv.thirdparty.deformable_attention.functions import ms_deform_attn_func
|
|
pre_compile_megatron_util()
|
|
# extension for all platform.
|
|
pre_compile_megatron_util()
|
|
|
|
|
|
if __name__ == '__main__':
|
|
pre_compile_all()
|