trending/miles
1
0
Fork 0
mirror of https://github.com/radixark/miles.git synced 2026-11-04 08:25:29 +08:00
Code Issues Packages Projects Releases Wiki Activity
Files
dfd65d5a8bd5d550d4e8a3ecb7abccfcce1eb4d2
miles/docker/npu_patch/miles.patch
T

564 lines
23 KiB
Diff
Raw Blame History

diff --git a/miles/backends/fsdp_utils/update_weight_utils.py b/miles/backends/fsdp_utils/update_weight_utils.py
index 98000e7d3..94fb71da5 100644
--- a/miles/backends/fsdp_utils/update_weight_utils.py
+++ b/miles/backends/fsdp_utils/update_weight_utils.py
@@ -234,12 +234,12 @@ class UpdateWeightFromDistributed(UpdateWeight):
i * self.args.rollout_num_gpus_per_engine + 1,
world_size,
self._group_name,
- backend="nccl",
+ backend="hccl",
)
for i, engine in enumerate(self.rollout_engines)
]
self._model_update_groups = init_process_group(
- backend="nccl",
+ backend="hccl",
init_method=f"tcp://{master_address}:{master_port}",
world_size=world_size,
rank=0,
diff --git a/miles/backends/megatron_utils/__init__.py b/miles/backends/megatron_utils/__init__.py
index 879b689c8..2f30a4bbd 100644
--- a/miles/backends/megatron_utils/__init__.py
+++ b/miles/backends/megatron_utils/__init__.py
@@ -1,6 +1,7 @@
import logging
import torch
+import mindspeed.megatron_adaptor
try:
import deep_ep
diff --git a/miles/backends/megatron_utils/actor.py b/miles/backends/megatron_utils/actor.py
index 54ca4b81a..3f79e49b7 100644
--- a/miles/backends/megatron_utils/actor.py
+++ b/miles/backends/megatron_utils/actor.py
@@ -39,12 +39,14 @@ from .parallel import verify_megatron_parallel_state
from .replay_utils import get_register_replay_list_func
from .update_weight.common import named_params_and_buffers
from .update_weight.update_weight_from_distributed.broadcast import UpdateWeightFromDistributed
-from .update_weight.update_weight_from_distributed.p2p import UpdateWeightP2P
+# from .update_weight.update_weight_from_distributed.p2p import UpdateWeightP2P
from .update_weight.update_weight_from_tensor import UpdateWeightFromTensor
logging.getLogger("megatron").setLevel(logging.WARNING)
logger = logging.getLogger(__name__)
+import mindspeed.megatron_adaptor
+from mindspeed.megatron_adaptor import repatch
class MegatronTrainRayActor(TrainRayActor):
@@ -58,7 +60,7 @@ class MegatronTrainRayActor(TrainRayActor):
monkey_patch_torch_dist()
super().init(args, role, with_ref)
-
+ repatch(args)
init(args)
if args.dumper_enable:
@@ -598,7 +600,7 @@ class MegatronTrainRayActor(TrainRayActor):
group_name = "actor_critic"
world_size = 2
self._actor_critic_groups = init_process_group(
- backend="nccl",
+ backend="hccl",
init_method=f"tcp://{master_address}:{master_port}",
world_size=world_size,
rank=0 if self.role == "actor" else 1,
diff --git a/miles/backends/megatron_utils/initialize.py b/miles/backends/megatron_utils/initialize.py
index dd3415c0a..15527171d 100644
--- a/miles/backends/megatron_utils/initialize.py
+++ b/miles/backends/megatron_utils/initialize.py
@@ -13,6 +13,7 @@ from miles.backends.training_utils.parallel import get_parallel_state, set_paral
from .parallel import create_megatron_parallel_state
logger = logging.getLogger(__name__)
+import mindspeed.megatron_adaptor
def _set_random_seed(
diff --git a/miles/backends/megatron_utils/model.py b/miles/backends/megatron_utils/model.py
index 5d2503dab..69051ddc5 100644
--- a/miles/backends/megatron_utils/model.py
+++ b/miles/backends/megatron_utils/model.py
@@ -142,6 +142,11 @@ def setup_model_and_optimizer(
kwargs[f.name] = getattr(args, f.name)
config = OptimizerConfig(**kwargs)
config.timers = None
+
+ for model_item in model:
+ for name, param in model_item.named_parameters():
+ if not param.is_leaf:
+ print(f"ERROR: param {name} is non-leaf, grad_fn={param.grad_fn}")
optimizer = get_megatron_optimizer(
config=config,
diff --git a/miles/backends/megatron_utils/model_provider.py b/miles/backends/megatron_utils/model_provider.py
index 45adc8679..c6c307a11 100644
--- a/miles/backends/megatron_utils/model_provider.py
+++ b/miles/backends/megatron_utils/model_provider.py
@@ -110,6 +110,20 @@ def get_model_provider_func(
provider.moe_router_bias_update_rate = args.moe_router_bias_update_rate
if getattr(args, "moe_aux_loss_coeff", None) is not None:
provider.moe_aux_loss_coeff = args.moe_aux_loss_coeff
+
+ if getattr(args, "gradient_accumulation_fusion", None) is not None:
+ provider.gradient_accumulation_fusion = args.gradient_accumulation_fusion
+
+ provider.moe_permute_fusion = args.moe_permute_fusion
+ provider.moe_aux_loss_coeff = args.moe_aux_loss_coeff
+
+ provider.freeze_language_model = False
+ provider.freeze_vision_model = False
+
+ for key, value in vars(args).items():
+ if hasattr(provider, key):
+ continue
+ setattr(provider, key, value)
provider.finalize()
def wrapped_bridge_provider(
diff --git a/miles/backends/megatron_utils/update_weight/common.py b/miles/backends/megatron_utils/update_weight/common.py
index 677f93c62..79ee57664 100644
--- a/miles/backends/megatron_utils/update_weight/common.py
+++ b/miles/backends/megatron_utils/update_weight/common.py
@@ -154,8 +154,8 @@ def named_params_and_buffers(
def _maybe_get_cpu_backup(x: torch.Tensor):
from torch_memory_saver import torch_memory_saver
- if (cpu_tensor := torch_memory_saver.get_cpu_backup(x)) is not None:
- return cpu_tensor
+ # if (cpu_tensor := torch_memory_saver.get_cpu_backup(x)) is not None:
+ # return cpu_tensor
return x
diff --git a/miles/backends/megatron_utils/update_weight/update_weight_from_distributed/broadcast.py b/miles/backends/megatron_utils/update_weight/update_weight_from_distributed/broadcast.py
index 51afedac4..c31c04a46 100644
--- a/miles/backends/megatron_utils/update_weight/update_weight_from_distributed/broadcast.py
+++ b/miles/backends/megatron_utils/update_weight/update_weight_from_distributed/broadcast.py
@@ -129,12 +129,12 @@ def connect_rollout_engines_from_distributed(
rank_cursor,
world_size,
group_name,
- backend="nccl",
+ backend="hccl",
)
)
rank_cursor += engine_gpu_counts[i]
model_update_groups = init_process_group(
- backend="nccl",
+ backend="hccl",
init_method=f"tcp://{master_address}:{master_port}",
world_size=world_size,
rank=0,
diff --git a/miles/ray/actor_group.py b/miles/ray/actor_group.py
index 669988b52..723113a69 100644
--- a/miles/ray/actor_group.py
+++ b/miles/ray/actor_group.py
@@ -85,7 +85,9 @@ class RayTrainGroup:
actor_impl = FSDPTrainRayActor
- TrainRayActor = ray.remote(num_gpus=1, runtime_env={"env_vars": env_vars})(actor_impl)
+ # TrainRayActor = ray.remote(num_gpus=1, runtime_env={"env_vars": env_vars})(actor_impl)
+ TrainRayActor = ray.remote(runtime_env={"env_vars": env_vars})(actor_impl)
+ device_name = "NPU"
# Create worker actors
actor_handles = []
@@ -93,11 +95,12 @@ class RayTrainGroup:
for rank in range(world_size):
actor = TrainRayActor.options(
num_cpus=num_gpus_per_actor,
- num_gpus=num_gpus_per_actor,
+ # num_gpus=num_gpus_per_actor,
scheduling_strategy=PlacementGroupSchedulingStrategy(
placement_group=pg,
placement_group_bundle_index=reordered_bundle_indices[rank],
),
+ resources={device_name: num_gpus_per_actor}
).remote(world_size, rank, master_addr, master_port)
if rank == 0:
master_addr, master_port = ray.get(actor.get_master_addr_and_port.remote())
diff --git a/miles/ray/placement_group.py b/miles/ray/placement_group.py
index bf8244b90..12442d9d6 100644
--- a/miles/ray/placement_group.py
+++ b/miles/ray/placement_group.py
@@ -14,10 +14,12 @@ from .rollout import RolloutManager
logger = logging.getLogger(__name__)
-@ray.remote(num_gpus=1)
+# @ray.remote(num_gpus=1)
+@ray.remote
class InfoActor:
def get_ip_and_gpu_id(self):
- return ray.util.get_node_ip_address(), ray.get_gpu_ids()[0]
+ # return ray.util.get_node_ip_address(), ray.get_gpu_ids()[0]
+ return ray.util.get_node_ip_address(), ray.get_runtime_context().get_accelerator_ids()["NPU"][0]
def sort_key(x):
@@ -38,12 +40,15 @@ def sort_key(x):
# representation that allows for sorting.
node_ip_parts = [ord(c) for c in node_identifier]
- return (node_ip_parts, gpu_id)
+ # return (node_ip_parts, gpu_id)
+ return (node_ip_parts, int(gpu_id))
def _create_placement_group(num_gpus):
"""Create a placement group with the specified number of GPUs."""
- bundles = [{"GPU": 1, "CPU": 1} for _ in range(num_gpus)]
+ # bundles = [{"GPU": 1, "CPU": 1} for _ in range(num_gpus)]
+ device_name = "NPU"
+ bundles = [{device_name: 1, "CPU": 1} for _ in range(num_gpus)]
pg = placement_group(bundles, strategy="PACK")
num_bundles = len(bundles)
@@ -56,7 +61,8 @@ def _create_placement_group(num_gpus):
scheduling_strategy=PlacementGroupSchedulingStrategy(
placement_group=pg,
placement_group_bundle_index=i,
- )
+ ),
+ resources={device_name: 1}
).remote()
)
gpu_ids = ray.get([actor.get_ip_and_gpu_id.remote() for actor in info_actors])
@@ -174,8 +180,9 @@ async def create_training_models(args, pgs, rollout_manager):
def create_rollout_manager(args, pg):
+ device_name = "NPU"
rollout_manager = RolloutManager.options(
- num_cpus=1, num_gpus=0, **(compute_ray_pin_head_options() if args.pin_rollout_manager_to_head else {})
+ num_cpus=1, resources={device_name: 0}, **(compute_ray_pin_head_options() if args.pin_rollout_manager_to_head else {})
).remote(args, pg)
# calculate num_rollout from num_epoch
diff --git a/miles/ray/rollout.py b/miles/ray/rollout.py
index ded71b511..89f515c9d 100644
--- a/miles/ray/rollout.py
+++ b/miles/ray/rollout.py
@@ -140,14 +140,15 @@ class ServerGroup:
}.items()
}
env_vars.update(dumper_utils.get_sglang_env(self.args))
-
+ device_name = "NPU"
rollout_engine = RolloutRayActor.options(
num_cpus=num_cpus,
- num_gpus=num_gpus,
+ # num_gpus=num_gpus,
scheduling_strategy=scheduling_strategy,
runtime_env={
"env_vars": env_vars,
},
+ resources={device_name: num_gpus}
).remote(
self.args,
rank=global_rank,
@@ -371,7 +372,9 @@ class RolloutManager:
init_http_client(args)
self.servers = start_rollout_servers(args, pg)
_start_session_server(args)
- self.rollout_engine_lock = Lock.options(num_cpus=1, num_gpus=0).remote()
+ # self.rollout_engine_lock = Lock.options(num_cpus=1, num_gpus=0).remote()
+ device_name = "NPU"
+ self.rollout_engine_lock = Lock.options(num_cpus=1, resources={device_name: 0}).remote()
self.rollout_id = -1
self._metric_checker = MetricChecker.maybe_create(args)
diff --git a/miles/ray/train_actor.py b/miles/ray/train_actor.py
index a4145ca27..33340323e 100644
--- a/miles/ray/train_actor.py
+++ b/miles/ray/train_actor.py
@@ -20,10 +20,17 @@ logger = logging.getLogger(__name__)
def get_local_gpu_id():
cvd = os.environ.get("CUDA_VISIBLE_DEVICES", None)
+ # if cvd is None:
+ # return ray.get_gpu_ids()[0]
+ # else:
+ # return cvd.split(",").index(str(ray.get_gpu_ids()[0]))
+ env_var = "ASCEND_RT_VISIBLE_DEVICES"
+ device_ids = ray.get_runtime_context().get_accelerator_ids()["NPU"]
+ cvd = os.environ.get(env_var, None)
if cvd is None:
- return ray.get_gpu_ids()[0]
+ return device_ids[0]
else:
- return cvd.split(",").index(str(ray.get_gpu_ids()[0]))
+ return cvd.split(",").index(str(device_ids[0]))
class TrainRayActor(RayActor):
diff --git a/miles/utils/arguments.py b/miles/utils/arguments.py
index abfa137ac..7c5fe007a 100644
--- a/miles/utils/arguments.py
+++ b/miles/utils/arguments.py
@@ -119,7 +119,7 @@ def get_miles_extra_args_provider(add_custom_arguments=None):
),
)
- reset_arg(parser, "--distributed-backend", type=str, default="nccl")
+ reset_arg(parser, "--distributed-backend", type=str, default="hccl")
reset_arg(parser, "--distributed-timeout-minutes", type=int, default=10)
return parser
@@ -2252,7 +2252,7 @@ def hf_validate_args(args, hf_config):
("tie_word_embeddings", "untie_embeddings_and_output_weights", lambda x, y: not x == y),
(
"rms_norm_eps",
- "norm_epsilon" if os.getenv("DEPRECATED_MEGATRON_COMPATIBLE", "0") == "1" else "layernorm_epsilon",
+ "norm_epsilon" if os.getenv("DEPRECATED_MEGATRON_COMPATIBLE", "0") == "1" else "norm_epsilon",
equal,
),
("rope_theta", "rotary_base", equal),
diff --git a/miles/utils/debug_utils/run_megatron/worker/main.py b/miles/utils/debug_utils/run_megatron/worker/main.py
index 359a51c1d..78e697e7c 100644
--- a/miles/utils/debug_utils/run_megatron/worker/main.py
+++ b/miles/utils/debug_utils/run_megatron/worker/main.py
@@ -117,9 +117,9 @@ def _parse_args() -> tuple[argparse.Namespace, WorkerScriptArgs]:
def _initialize_megatron(args: argparse.Namespace) -> None:
- torch.distributed.init_process_group(backend="nccl")
+ torch.distributed.init_process_group(backend="hccl")
local_rank: int = int(os.environ.get("LOCAL_RANK", 0))
- torch.cuda.set_device(local_rank)
+ torch.npu.set_device(local_rank)
args.hf_checkpoint = str(args.script_hf_checkpoint)
args.__dict__.setdefault("megatron_to_hf_mode", "raw")
diff --git a/miles/utils/external_utils/command_utils.py b/miles/utils/external_utils/command_utils.py
index d016e01ac..08b4d6eff 100644
--- a/miles/utils/external_utils/command_utils.py
+++ b/miles/utils/external_utils/command_utils.py
@@ -193,6 +193,107 @@ def execute_train(
)
+def execute_train_npu(
+ train_args: str,
+ num_gpus_per_node: int,
+ megatron_model_type: str | None,
+ train_script: str = "train.py",
+ before_ray_job_submit=None,
+ extra_env_vars=None,
+ config: ExecuteTrainConfig | None = None,
+ megatron_path: str = "/root/Megatron-LM",
+):
+ if extra_env_vars is None:
+ extra_env_vars = {}
+ if config is None:
+ config = ExecuteTrainConfig()
+ external_ray = get_bool_env_var("SLIME_SCRIPT_EXTERNAL_RAY")
+ master_addr = os.environ.get("MASTER_ADDR", "127.0.0.1")
+
+ train_backend_fsdp = "--train-backend fsdp" in train_args
+ assert train_backend_fsdp == (megatron_model_type is None)
+
+ exec_command_cpu(
+ "pkill -9 sglang; "
+ "sleep 3; "
+ f"{'' if external_ray else 'ray stop --force; '}"
+ f"{'' if external_ray else 'pkill -9 ray; '}"
+ # cannot be run in CI, o/w kill the parent script
+ # "pkill -9 python; "
+ "pkill -9 slime; "
+ "sleep 3; "
+ f"{'' if external_ray else 'pkill -9 ray; '}"
+ # "pkill -9 python; "
+ "pkill -9 slime; "
+ "pkill -9 redis; "
+ "true; "
+ )
+
+ if not external_ray:
+ exec_command_cpu(
+ # will prevent ray from buffering stdout/stderr
+ f"export PYTHONUNBUFFERED=1 && "
+ f"ray start --head --node-ip-address {master_addr} --disable-usage-stats "
+ )
+
+ f = before_ray_job_submit
+ if f is not None:
+ f()
+
+ runtime_env_json = json.dumps(
+ {
+ "env_vars": {
+ "CUDA_DEVICE_MAX_CONNECTIONS": "1",
+ "RAY_EXPERIMENTAL_NOSET_ASCEND_RT_VISIBLE_DEVICES": "1",
+ # Replace with actual Ascend toolkit paths
+ "ASCEND_TOOLKIT_HOME": "/home/w00664509/ascend-toolkit/latest/",
+ "ASCEND_OPP_PATH": "/home/w00664509/cann8.5.2/ascend-toolkit/latest/opp/",
+ "ASCEND_AICPU_PATH": "/home/w00664509/cann8.5.2/ascend-toolkit/latest/",
+ "ASCEND_HOME_PATH": "/home/w00664509/cann8.5.2/ascend-toolkit/latest/",
+ "set_env_path": "/home/w00664509/cann8.5.2/nnal/atb/set_env.sh",
+ "HYDRA_FULL_ERROR": "1",
+ "HCCL_HOST_SOCKET_PORT_RANGE": "60000-60050",
+ "HCCL_NPU_SOCKET_PORT_RANGE": "61000-61050",
+ # If setting this in FSDP, the computation communication overlapping may have issues
+ **(
+ {}
+ if train_backend_fsdp
+ else {
+ "CUDA_DEVICE_MAX_CONNECTIONS": "1",
+ }
+ ),
+ "NCCL_NVLS_ENABLE": str(int(check_has_nvlink())),
+ "no_proxy": f"127.0.0.1,{master_addr}",
+ # This is needed by megatron / torch distributed in multi-node setup
+ "MASTER_ADDR": master_addr,
+ **(
+ {
+ "CUDA_ENABLE_COREDUMP_ON_EXCEPTION": "1",
+ "CUDA_COREDUMP_SHOW_PROGRESS": "1",
+ "CUDA_COREDUMP_GENERATION_FLAGS": "skip_nonrelocated_elf_images,skip_global_memory,skip_shared_memory,skip_local_memory,skip_constbank_memory",
+ "CUDA_COREDUMP_FILE": "/root/shared_data/cuda_coredump_%h.%p.%t",
+ }
+ if config.cuda_core_dump
+ else {}
+ ),
+ **extra_env_vars,
+ **_parse_extra_env_vars(config.extra_env_vars),
+ }
+ }
+ )
+
+ if get_bool_env_var("SLIME_SCRIPT_ENABLE_RAY_SUBMIT", "1"):
+ model_args = shell_safe_model_args(megatron_model_type)
+ exec_command_cpu(
+ f"export no_proxy=127.0.0.1 && export PYTHONUNBUFFERED=1 && "
+ f'ray job submit --address="http://127.0.0.1:8265" '
+ f"--runtime-env-json='{runtime_env_json}' "
+ f"-- python3 {train_script} "
+ f"{model_args} "
+ f"{train_args}"
+ )
+
+
def _parse_extra_env_vars(text: str):
try:
return json.loads(text)
diff --git a/miles/utils/ppo_utils.py b/miles/utils/ppo_utils.py
index 4c09dc508..369d8fda5 100644
--- a/miles/utils/ppo_utils.py
+++ b/miles/utils/ppo_utils.py
@@ -10,7 +10,7 @@ import torch.nn.functional as F
from miles.backends.training_utils.parallel import get_parallel_state
-@torch.compile(dynamic=True)
+# @torch.compile(dynamic=True)
def compute_approx_kl(
log_probs: torch.Tensor,
log_probs_base: torch.Tensor,
@@ -123,7 +123,7 @@ def compute_gspo_kl(
return ppo_kl
-@torch.compile(dynamic=True)
+# @torch.compile(dynamic=True)
def compute_policy_loss(
ppo_kl: torch.Tensor,
advantages: torch.Tensor,
@@ -166,7 +166,7 @@ class _VocabParallelEntropy(torch.autograd.Function):
@staticmethod
def forward(ctx, vocab_parallel_logits: torch.Tensor, process_group: dist.ProcessGroup) -> torch.Tensor:
- @torch.compile(dynamic=True)
+ # @torch.compile(dynamic=True)
def mul_reduce(a, b):
return (a * b).sum(dim=-1, keepdim=True)
diff --git a/miles/utils/reloadable_process_group.py b/miles/utils/reloadable_process_group.py
index af4878b4e..f796a4ee7 100644
--- a/miles/utils/reloadable_process_group.py
+++ b/miles/utils/reloadable_process_group.py
@@ -154,7 +154,7 @@ class ReloadableProcessGroup(torch.distributed.ProcessGroup):
for reloadable_group in reloadable_groups:
if reloadable_group.group is not None:
continue
- group = old_new_group(ranks=reloadable_group.group_info["ranks"], backend="nccl")
+ group = old_new_group(ranks=reloadable_group.group_info["ranks"], backend="hccl")
reloadable_group.group = group
def rank(self) -> int:
diff --git a/miles/utils/tensor_backper.py b/miles/utils/tensor_backper.py
index 2fc2a6359..6952d1fc1 100644
--- a/miles/utils/tensor_backper.py
+++ b/miles/utils/tensor_backper.py
@@ -110,6 +110,6 @@ def _compute_hash_tensor(x: torch.Tensor):
# Not a real/good hash, but pretty fast
x = x.contiguous()
x = x.view(-1)
- x = x.view(torch.uint32)
+ x = x.view(torch.int32)
x = x.sum()
return x.item()
diff --git a/tools/convert_hf_to_torch_dist.py b/tools/convert_hf_to_torch_dist.py
index 250af0b68..8e407b471 100644
--- a/tools/convert_hf_to_torch_dist.py
+++ b/tools/convert_hf_to_torch_dist.py
@@ -17,6 +17,7 @@ from miles.backends.megatron_utils.model_provider import get_model_provider_func
from miles.utils.logging_utils import configure_logger
from miles.utils.memory_utils import print_memory
from miles_plugins.models.hf_attention import _load_hf_config
+import mindspeed.megatron_adaptor
def add_convertion_args(parser):
@@ -94,17 +95,17 @@ def main():
local_rank = int(os.getenv("LOCAL_RANK") or os.getenv("SLURM_LOCALID") or 0)
global_rank = int(os.getenv("RANK") or os.getenv("SLURM_PROCID") or 0)
- torch.cuda.set_device(local_rank)
+ torch.npu.set_device(local_rank)
os.environ.setdefault("WORLD_SIZE", str(world_size))
os.environ.setdefault("RANK", str(global_rank))
os.environ.setdefault("LOCAL_RANK", str(local_rank))
os.environ.setdefault("MASTER_ADDR", "localhost")
os.environ.setdefault("MASTER_PORT", "12355")
dist.init_process_group(
- backend="nccl",
+ backend="hccl",
world_size=world_size,
rank=global_rank,
- device_id=torch.device(f"cuda:{local_rank}"),
+ device_id=torch.device(f"npu:{local_rank}"),
)
args = get_args()
init(args)
@@ -122,9 +123,9 @@ def main():
print(f"Model loaded: {hf_model_path}")
print_memory("after loading model")
- torch.cuda.synchronize()
+ torch.npu.synchronize()
gc.collect()
- torch.cuda.empty_cache()
+ torch.npu.empty_cache()
save_checkpoint(1, model, None, None, 0)
diff --git a/train.py b/train.py
index 9a4a29550..0512bee06 100644
--- a/train.py
+++ b/train.py
@@ -8,6 +8,7 @@ from miles.utils.async_utils import eager_create_task
from miles.utils.logging_utils import configure_logger
from miles.utils.misc import should_run_periodic_action
from miles.utils.tracking_utils import finish_tracking, init_tracking
+import mindspeed.megatron_adaptor
async def train(args):
Reference in New Issue View Git Blame Copy Permalink