eplb 采集-MindSpeed-RL-AtomGit

Mcli/eplb_generate_map_ds.py+32-34

	@@ -20,6 +20,7 @@ def find_replacement(current_group, expert_idx, weight, used_experts):
20		20
21		21
22	def tensor_to_json(tensor: torch.Tensor, num_gpus: int, output_path):	22	def tensor_to_json(tensor: torch.Tensor, num_gpus: int, output_path):
		23	+ # 数据格式转换 Tensor -> Json 目标格式
23	num_layers, num_replicas = tensor.shape	24	num_layers, num_replicas = tensor.shape
24	experts_per_device = num_replicas // num_gpus	25	experts_per_device = num_replicas // num_gpus
25		26
	@@ -58,38 +59,38 @@ def tensor_to_json(tensor: torch.Tensor, num_gpus: int, output_path):
58		59
59		60
60	def load_and_aggregate(json_folder: str) -> torch.Tensor:	61	def load_and_aggregate(json_folder: str) -> torch.Tensor:
61	- # 构造单一文件路径	62	+ # 读取并统计所有热点数据
62	- filepath = os.path.join(json_folder, "token_collects_all.json")	63	+ files: List[str] = [
63	- if not os.path.exists(filepath):	64	+ os.path.join(json_folder, f)
64	- raise FileNotFoundError(f"File not found: {filepath}")	65	+ for f in os.listdir(json_folder)
65	-	66	+ if f.endswith(".json")
66	- # 读取并规范化键为 int，同时统计最大 layer / expert	67	+ ]
67	- with open(filepath, "r") as f:	68	+ if not files:
68	- raw = json.load(f)	69	+ raise FileNotFoundError(f"No .json files found in: {json_folder}")
69	-	70	+ loaded: List[Dict[int, Dict[int, int]]] = []
70	- norm: Dict[int, Dict[int, int]] = {}
71	max_layer = -1	71	max_layer = -1
72	max_expert = -1	72	max_expert = -1
73	-	73	+ for path in files:
74	- for layer_id_str, experts_map in raw.items():	74	+ with open(path, "r") as f:
75	- layer_id = int(layer_id_str)	75	+ raw = json.load(f)
76	- expert_dict = {int(eid): int(cnt) for eid, cnt in experts_map.items()}	76	+ norm: Dict[int, Dict[int, int]] = {}
77	- norm[layer_id] = expert_dict	77	+ for layer_id_str, experts_map in raw.items():
78	-	78	+ layer_id = int(layer_id_str)
79	- if expert_dict:	79	+ expert_dict = {int(eid): int(cnt) for eid, cnt in experts_map.items()}
80	- max_expert = max(max_expert, max(expert_dict.keys()))	80	+ norm[layer_id] = expert_dict
81	- max_layer = max(max_layer, layer_id)	81	+ if expert_dict:
82	-	82	+ max_expert = max(max_expert, max(expert_dict.keys()))
		83	+ max_layer = max(max_layer, layer_id)
		84	+ loaded.append(norm)
83	# 初始化 [num_layers, num_experts] 矩阵	85	# 初始化 [num_layers, num_experts] 矩阵
84	num_layers = max_layer + 1 if max_layer >= 0 else 0	86	num_layers = max_layer + 1 if max_layer >= 0 else 0
85	num_experts = max_expert + 1 if max_expert >= 0 else 0	87	num_experts = max_expert + 1 if max_expert >= 0 else 0
86	mat = torch.zeros((num_layers, num_experts), dtype=torch.long)	88	mat = torch.zeros((num_layers, num_experts), dtype=torch.long)
87	-	89	+ # 累加
88	- # 累加到矩阵	90	+ for data in loaded:
89	- for layer_id, experts in norm.items():	91	+ for layer_id, experts in data.items():
90	- for eid, cnt in experts.items():	92	+ for eid, cnt in experts.items():
91	- mat[layer_id, eid] += cnt	93	+ mat[layer_id, eid] += cnt
92	-
93	return mat	94	return mat
94		95
95		96
	@@ -109,23 +110,21 @@ if __name__ == "__main__":
109	num_nodes = args.num_nodes	110	num_nodes = args.num_nodes
110	num_gpus = args.num_gpus	111	num_gpus = args.num_gpus
111	output_path = args.output_path	112	output_path = args.output_path
112	-	113	+ # 计算冗余专家
113	phy2log, log2phy, logcnt = eplb.rebalance_experts(weight, num_replicas, num_groups, num_nodes, num_gpus)	114	phy2log, log2phy, logcnt = eplb.rebalance_experts(weight, num_replicas, num_groups, num_nodes, num_gpus)
114	- #一张卡上的冗余专家数
115	experts_per_gpu = num_replicas // num_gpus	115	experts_per_gpu = num_replicas // num_gpus
116	- #总本地专家数
117	num_experts = weight.size(1)	116	num_experts = weight.size(1)
118	-	117	+ # 冗余专家去重
119	for layer in range(weight.size(0)):	118	for layer in range(weight.size(0)):
120	weight_layer = weight[layer]	119	weight_layer = weight[layer]
121	for gpu in range(num_gpus):	120	for gpu in range(num_gpus):
122	start = gpu * experts_per_gpu	121	start = gpu * experts_per_gpu
123	end = start + experts_per_gpu	122	end = start + experts_per_gpu
124	- #遍历当前层当前gpu上的专家列表	123	+ # 遍历当前层的专家列表
125	group = phy2log[layer, start:end]	124	group = phy2log[layer, start:end]
126	- #用于记录当前组中已经出现过的专家	125	+ # 用于记录当前组中已经出现过的专家
127	seen = set()	126	seen = set()
128	- #用于记录已经尝试替换但不合适的专家，避免重复尝试	127	+ # 用于记录已经尝试替换但不合适的专家
129	used_experts = set()	128	used_experts = set()
130	for j in range(experts_per_gpu):	129	for j in range(experts_per_gpu):
131	expert = group[j].item()	130	expert = group[j].item()
	@@ -136,5 +135,4 @@ if __name__ == "__main__":
136	replacement = find_replacement(group.tolist(), expert, weight_layer, used_experts)	135	replacement = find_replacement(group.tolist(), expert, weight_layer, used_experts)
137	phy2log[layer, start + j] = replacement	136	phy2log[layer, start + j] = replacement
138	seen.add(phy2log[layer, start + j].item())	137	seen.add(phy2log[layer, start + j].item())
139	-
140	tensor_to_json(phy2log, num_gpus, output_path)	138	tensor_to_json(phy2log, num_gpus, output_path)

Mdocs/features/EPLB.md+28-11

	@@ -31,7 +31,7 @@
31		31
32	* 在 MoE 前向计算时，系统会拦截 `AscendUnquantizedFusedMoEMethod.apply` 调用。	32	* 在 MoE 前向计算时，系统会拦截 `AscendUnquantizedFusedMoEMethod.apply` 调用。
33	* 自动统计 `topk_ids`（即每个 token 被路由到的 expert ID），并按 layer 累计。	33	* 自动统计 `topk_ids`（即每个 token 被路由到的 expert ID），并按 layer 累计。
34	-* 所有数据会收集到 rank0 所在的节点上，统一写入一个 JSON 文件。	34	+* 每个rank的数据会收集各自所在的节点上，并保存各自对应的 JSON 文件。
35	* JSON 文件格式如下：	35	* JSON 文件格式如下：
36		36
37	```json	37	```json
	@@ -56,6 +56,7 @@
56	* 内层 key 表示 expert ID	56	* 内层 key 表示 expert ID
57	* 值为该 expert 被分配的 token 数	57	* 值为该 expert 被分配的 token 数
58		58
		59	+* 收集每个节点上的json文件到主节点上
59	---	60	---
60		61
61	### 2. 配置方法	62	### 2. 配置方法
	@@ -65,7 +66,7 @@
65	```yaml	66	```yaml
66	generate:	67	generate:
67	token_collects: true # 是否开启 token 收集	68	token_collects: true # 是否开启 token 收集
68	- token_save_path: "/path/to/save" # JSON 文件保存目录	69	+ token_save_path: "/path/to/save" # JSON 文件保存目录, 默认为json_file
69	```	70	```
70		71
71	参数说明：	72	参数说明：
	@@ -77,33 +78,49 @@ generate:
77	* `token_save_path`:	78	* `token_save_path`:
78		79
79	* 必填，指定 JSON 文件存储目录	80	* 必填，指定 JSON 文件存储目录
80	- * 最终只会生成一个汇总文件：	81	+ * 最终每个rank收集到的数据会保存在其所在节点上：
81		82
82	```	83	```
83	- token_collects_all.json	84	+ "eplb_token_collects_{rank}.json"
84	```	85	```
85		86
		87	+
		88	+在 `examples\eplb\collect_json_file.sh` 中修改以下参数：
		89	+```yaml
		90	+# 远程服务器列表（格式：user@ip）
		91	+SERVERS = (
		92	+ "root@IP"
		93	+ "root@IP"
		94	+ )
		95	+```
		96	+参数说明：
		97	+
		98	+* `SERVERS`:
		99	+
		100	+ * 必填，指定远程服务器列表
		101	+
86	---	102	---
87		103
88	### 3. 运行效果	104	### 3. 运行效果
89		105
90	-1. 运行推理或训练后，各 rank 会先本地统计数据，再通过分布式通信汇总到 rank0。	106	+1. 运行推理或训练后，各 rank 会本地统计数据，并存储在当前节点的`token_save_path`目录下。
91		107
92	-2. rank0 会在 `token_save_path` 下生成一个 JSON 文件：	108	+2. `token_save_path` 下会生成多个 JSON 文件：
93		109
94	```bash	110	```bash
95	- $ ls /mnt/data2/MindSpeed-RL	111	+ eplb_token_collects_{rank}.json
96	- token_collects_all.json
97	```	112	```
98		113
99	-3. 文件中保存了所有 rank 的统计结果，包含每层每个 expert 的 token 数。	114	+3. 文件中保存了各个 rank 的统计结果，包含每层每个 expert 的 token 数。
		115	+
		116	+4. 在主节点运行`examples\eplb\collect_json_file.sh`后，各个节点上的 JSON 文件会被统一收集到主节点上。
100		117
101	---	118	---
102		119
103	### 4. 注意事项	120	### 4. 注意事项
104		121
105	* 性能开销：开启采集后会增加统计与 JSON I/O，建议在 profiling 或 debug 时使用。	122	* 性能开销：开启采集后会增加统计与 JSON I/O，建议在 profiling 或 debug 时使用。
106	-* 多机环境：数据会自动聚合，用户只需查看 rank0 生成的 `token_collects_all.json` 文件。	123	+* 多机环境：共享存储场景下，不需要在主节点运行`examples\eplb\collect_json_file.sh` 所有 JSON 文件会被统一收集到共享存储中。
107	* 文件写入安全：采用临时文件 + 原子替换，避免因进程异常退出导致 JSON 文件损坏。	124	* 文件写入安全：采用临时文件 + 原子替换，避免因进程异常退出导致 JSON 文件损坏。
108	* 配置优先级：`yaml` 中的配置会覆盖代码默认值。	125	* 配置优先级：`yaml` 中的配置会覆盖代码默认值。
109	---	126	---
	@@ -185,7 +202,7 @@ python mindspeed_rl/workers/eplb/eplb_generate_map_ds.py \
185		202
186	* 其中：	203	* 其中：
187		204
188	- * json_folder 表示 `token_collects_all.json` 所在目录，与 `Step1 MOE Token Collect` 中的 `token_save_path` 一致	205	+ * json_folder 表示主节点上 `eplb_token_collects_{rank}.json` 所在目录，与 `Step1 MOE Token Collect` 中的 `token_save_path` 一致
189	* num_replicas 表示冗余专家总数（含原始专家）	206	* num_replicas 表示冗余专家总数（含原始专家）
190	* num_groups 表示负载均衡策略中的专家分组数, 要求能被num_gpus整除	207	* num_groups 表示负载均衡策略中的专家分组数, 要求能被num_gpus整除
191	* num_nodes 表示机器数	208	* num_nodes 表示机器数

Aexamples/eplb/collect_json_file.sh+24-0

@@ -0,0 +1,24 @@
	1	+# 本地保存路径
	2	+LOCAL_DIR="./MindSpeed-RL/json_file"
	3	+
	4	+# 远程服务器列表（格式：user@ip）
	5	+SERVERS=(
	6	+ "root@IP"
	7	+ "root@IP"
	8	+)
	9	+
	10	+# 远程 json_file 文件夹路径
	11	+REMOTE_DIR="./MindSpeed-RL/json_file"
	12	+
	13	+# 遍历每台服务器并同步文件
	14	+for SERVER in "${SERVERS[@]}"; do
	15	+ scp "$SERVER:$REMOTE_DIR/*.json" "$LOCAL_DIR/"
	16	+ if [ $? -ne 0 ]; then
	17	+ echo " 从 $SERVER 同步失败，请检查 SSH 连接或路径是否正确"
	18	+ else
	19	+ echo " $SERVER 同步完成"
	20	+ fi
	21	+done
	22	+
	23	+echo " 所有文件已同步到 $LOCAL_DIR"
	24	+echo " 当前目录中 JSON 文件数量：$(ls -1 "$LOCAL_DIR"/*.json 2>/dev/null \| wc -l)"

Mexamples/eplb/grpo_trainer_deepseek_r1_671b_eplb.sh+1-0

	@@ -10,6 +10,7 @@ echo "Use $YAML"
10		10
11	ulimit -n 32768	11	ulimit -n 32768
12	mkdir logs	12	mkdir logs
		13	+mkdir json_file
13		14
14	export TASK_QUEUE_ENABLE=2	15	export TASK_QUEUE_ENABLE=2
15	export HCCL_IF_BASE_PORT=24703	16	export HCCL_IF_BASE_PORT=24703

Mmindspeed_rl/models/rollout/vllm_adapter/fused_moe.py+18-43

	@@ -251,7 +251,6 @@ def fused_experts_with_all2all(
251	num_tokens, _ = hidden_states.shape	251	num_tokens, _ = hidden_states.shape
252	num_experts = w1.shape[0]	252	num_experts = w1.shape[0]
253	device = hidden_states.device	253	device = hidden_states.device
254	-
255	# EPLB update global_num_experts	254	# EPLB update global_num_experts
256	if expert_map is not None:	255	if expert_map is not None:
257	global_num_experts = len(expert_map) + global_redundant_expert_num	256	global_num_experts = len(expert_map) + global_redundant_expert_num
	@@ -273,6 +272,7 @@ def fused_experts_with_all2all(
273		272
274	global_expert_tokens = torch.bincount(expanded_expert_idx,	273	global_expert_tokens = torch.bincount(expanded_expert_idx,
275	minlength=global_num_experts)	274	minlength=global_num_experts)
		275	+
276	scatter_sizes = global_expert_tokens.view(ep_group.world_size,	276	scatter_sizes = global_expert_tokens.view(ep_group.world_size,
277	-1).sum(-1)	277	-1).sum(-1)
278		278
	@@ -462,50 +462,25 @@ class AscendUnquantizedFusedMoEMethod(UnquantizedFusedMoEMethod):
462	eplb_token_collects, eplb_token_save_path = get_EPLB_args()	462	eplb_token_collects, eplb_token_save_path = get_EPLB_args()
463	if eplb_token_collects:	463	if eplb_token_collects:
464	self.layers = layer.moe_instance_id	464	self.layers = layer.moe_instance_id
465	- rank = dist.get_rank()	465	+ filepath = os.path.join(eplb_token_save_path, f"eplb_token_collects_{torch.distributed.get_rank()}.json")
466	- world_size = dist.get_world_size()
467		466
468	- # cur_rank	467	+ if os.path.exists(filepath):
469	- local_stats = {self.layers: {eid: 0 for eid in range(global_num_experts)}}	468	+ with open(filepath, "r") as f:
		469	+ experts_loads = json.load(f)
		470	+ # key 转 int
		471	+ experts_loads = {int(k): {int(e): v for e, v in vdict.items()} for k, vdict in experts_loads.items()}
		472	+ else:
		473	+ experts_loads = {}
		474	+
		475	+ # 当前层的统计
		476	+ if self.layers not in experts_loads:
		477	+ experts_loads[self.layers] = {eid: 0 for eid in range(global_num_experts)}
		478	+
470	unique_ids, counts = torch.unique(topk_ids, return_counts=True)	479	unique_ids, counts = torch.unique(topk_ids, return_counts=True)
471	- for uid, cnt in zip(unique_ids.tolist(), counts.tolist()):	480	+ for eid, cnt in zip(unique_ids.tolist(), counts.tolist()):
472	- local_stats[self.layers][uid] += cnt	481	+ experts_loads[self.layers][eid] += cnt
473	-	482	+ with open(filepath, "w") as f:
474	- all_stats = [None for _ in range(world_size)]	483	+ json.dump(experts_loads, f, indent=2)
		Y Y yezhibei2025年9月25日采集完成后在 eplb_token_save_path 路径下保存每个 rank 上收集的JSON文件 Z zhoubeirong 2025年9月26日评论： 2025年9月26日评论： Y yezhibei 2025年9月26日评论： 2025年9月26日评论：
475	- dist.all_gather_object(all_stats, local_stats)
476	-
477	- # all ranks
478	- if rank == 0:
479	- # merge
480	- merged_stats = {}
481	- for rank_stats in all_stats:
482	- for layer_id, e_dict in rank_stats.items():
483	- if layer_id not in merged_stats:
484	- merged_stats[layer_id] = {eid: 0 for eid in range(global_num_experts)}
485	- for eid, cnt in e_dict.items():
486	- merged_stats[layer_id][eid] += cnt
487	-
488	- filepath = os.path.join(eplb_token_save_path, f"token_collects_all.json")
489	- if os.path.exists(filepath):
490	- with open(filepath, "r") as f:
491	- old_stats = json.load(f)
492	- old_stats = {int(k): {int(e): v for e, v in vdict.items()} for k, vdict in old_stats.items()}
493	- else:
494	- old_stats = {}
495	-
496	- # merge old_stats
497	- for layer_id, e_dict in merged_stats.items():
498	- if layer_id not in old_stats:
499	- old_stats[layer_id] = {eid: 0 for eid in range(global_num_experts)}
500	- for eid, cnt in e_dict.items():
501	- old_stats[layer_id][eid] += int(cnt)
502	-
503	- with tempfile.NamedTemporaryFile("w", delete=False, dir=os.path.dirname(filepath)) as tmpfile:
504	- json.dump(old_stats, tmpfile, indent=2)
505	- tmpfile.flush()
506	- os.fsync(tmpfile.fileno())
507	- temp_name = tmpfile.name
508	- shutil.move(temp_name, filepath)
509		484
510	topk_weights = topk_weights.to(x.dtype)	485	topk_weights = topk_weights.to(x.dtype)
511		486