streamline code; add intermediate saving support for ep

This commit is contained in:
ZihanWang314
2025-05-22 07:21:52 +00:00
parent 98fd21ce21
commit f38f67706c
123 changed files with 710 additions and 5601 deletions

View File

@@ -30,9 +30,10 @@ def get_summary(files):
expert_ids, expert_weights = parse_line(line)
np.add.at(gate_scores[layer_id], expert_ids, expert_weights)
np.add.at(token_scores[layer_id], expert_ids, np.ones_like(expert_weights) / TOP_K)
gate_scores = gate_scores / np.sum(gate_scores, axis=0)
token_scores = token_scores / np.sum(token_scores, axis=0)
total = sum(token_scores[0])
gate_scores = gate_scores / total
token_scores = token_scores / total
summary = {"token_scores": token_scores, "gate_scores": gate_scores}
summary = {k: {str(i+1): {str(j): round(v, 4) for j, v in enumerate(l)} for i, l in enumerate(v)} for k, v in summary.items()}
@@ -65,7 +66,6 @@ if __name__ == "__main__":
for file in os.listdir(os.path.join(args.expert_scores_dir, rank)):
file_names.append([rank, file])
summary_file = os.path.join(args.expert_scores_dir, "summary.json")
summary = get_summary(file_names)