migrate-benchmarks-to-v2.12.0-pytorch-AtomGit

Abenchmarks/inference/DCNv2/dcnv2.patch+22-0

@@ -0,0 +1,22 @@
	1	+--- CTR_Algorithm/pytorch/model/DCNv2.py 2026-03-31 11:09:49.207077300 +0800
	2	++++ dcn/CTR_Algorithm/pytorch/model/DCNv2.py 2026-03-31 11:12:56.799077300 +0800
	3	+@@ -120,7 +120,8 @@
	4	+ """
	5	+ super(DeepCrossNetv2, self).__init__()
	6	+ self.feature_fields = feature_fields
	7	+- self.offsets = np.array((0, *np.cumsum(feature_fields)[:-1]), dtype = np.long)
	8	++ self.offsets = np.array((0, *np.cumsum(feature_fields)[:-1]), dtype = np.int64)
	9	++ self.register_buffer("offsets_tensor", torch.as_tensor(self.offsets, dtype=torch.int64))
	10	+ self.model_method = model_method
	11	+
	12	+ # Embedding layer
	13	+@@ -157,7 +158,7 @@
	14	+ raise NotImplementedError
	15	+
	16	+ def forward(self, x):
	17	+- tmp = x + x.new_tensor(self.offsets).unsqueeze(0)
	18	++ tmp = x + self.offsets_tensor.unsqueeze(0)
	19	+
	20	+ # embeded dense vector
	21	+ embeded_x = self.embedding(tmp).view(-1, self.embedding_out_dim)
	22	+

Abenchmarks/inference/DCNv2/dcnv2.py+205-0

@@ -0,0 +1,205 @@
	1	+# -- coding: utf-8 --
	2	+import os
	3	+import pandas as pd
	4	+import numpy as np
	5	+import torch
	6	+import torch.nn as nn
	7	+import torch.optim as optim
	8	+import torch.utils.data as Data
	9	+import torch.nn.functional as F
	10	+from sklearn.preprocessing import LabelEncoder
	11	+from sklearn.model_selection import train_test_split
	12	+import sys
	13	+import time
	14	+import argparse
	15	+
	16	+from CTR_Algorithm.data.AmazonDataPreprocess import AmazonBookPreprocess
	17	+
	18	+MODEL_NAME='DCNV2'
	19	+
	20	+data = pd.read_csv('CTR_Algorithm/data/data.csv')
	21	+data_X = data.iloc[:,2:]
	22	+data_y = data.click.values
	23	+
	24	+data_X = data_X.apply(LabelEncoder().fit_transform)
	25	+
	26	+fields = data_X.max().values + 1
	27	+
	28	+# data_X = data.iloc[:,:-1]
	29	+# data_y = data.label.values
	30	+
	31	+tmp_X, test_X, tmp_y, test_y = train_test_split(data_X, data_y, test_size = 0.2, random_state=42, stratify=data_y)
	32	+train_X, val_X, train_y, val_y = train_test_split(tmp_X, tmp_y, test_size = 0.25, random_state=42, stratify=tmp_y)
	33	+
	34	+train_X = torch.from_numpy(train_X.values).long()
	35	+val_X = torch.from_numpy(val_X.values).long()
	36	+test_X = torch.from_numpy(test_X.values).long()
	37	+
	38	+train_y = torch.from_numpy(train_y).long()
	39	+val_y = torch.from_numpy(val_y).long()
	40	+test_y = torch.from_numpy(test_y).long()
	41	+
	42	+train_set = Data.TensorDataset(train_X, train_y)
	43	+val_set = Data.TensorDataset(val_X, val_y)
	44	+train_loader = Data.DataLoader(dataset=train_set,
	45	+ batch_size=32,
	46	+ shuffle=True)
	47	+val_loader = Data.DataLoader(dataset=val_set,
	48	+ batch_size=32,
	49	+ shuffle=False)
	50	+
	51	+
	52	+def detect_device_type():
	53	+ try:
	54	+ import torch
	55	+ if torch.cuda.is_available():
	56	+ return "cuda"
	57	+ try:
	58	+ import torch_npu
	59	+ import os
	60	+ os.environ['TORCHINDUCTOR_NPU_BACKEND'] = 'mlir'
	61	+ if torch.npu.is_available():
	62	+ return "npu"
	63	+ except ImportError:
	64	+ pass
	65	+ except ImportError:
	66	+ pass
	67	+ return "cpu"
	68	+
	69	+
	70	+def get_profile(profiler_start_step: int, profiler_end_step: int, profiling_save_path: str):
	71	+ warm_step = profiler_start_step
	72	+ active_step = profiler_end_step - warm_step +1
	73	+ print(f"[Profile INFO]: warm_step: {warm_step}, active_step: {active_step}, profiling_save_path: {profiling_save_path}")
	74	+ device = detect_device_type()
	75	+ if device == 'npu':
	76	+ import torch_npu
	77	+ g_prof_config = torch_npu.profiler._ExperimentalConfig(
	78	+ export_type=[
	79	+ torch_npu.profiler.ExportType.Text,
	80	+ torch_npu.profiler.ExportType.Db
	81	+ ],
	82	+ profiler_level=torch_npu.profiler.ProfilerLevel.Level2,
	83	+ msprof_tx=False,
	84	+ aic_metrics=torch_npu.profiler.AiCMetrics.AiCoreNone,
	85	+ l2_cache=False,
	86	+ op_attr=False,
	87	+ data_simplification=False,
	88	+ record_op_args=False,
	89	+ gc_detect_threshold=None)
	90	+
	91	+ return torch_npu.profiler.profile(
	92	+ activities=[
	93	+ torch_npu.profiler.ProfilerActivity.CPU,
	94	+ torch_npu.profiler.ProfilerActivity.NPU],
	95	+ schedule=torch_npu.profiler.schedule(wait=0, warmup=warm_step, active=active_step, repeat=1),
	96	+ on_trace_ready=torch_npu.profiler.tensorboard_trace_handler(profiling_save_path),
	97	+ record_shapes=False,
	98	+ profile_memory=False,
	99	+ with_stack=False,
	100	+ with_modules=False,
	101	+ with_flops=False,
	102	+ experimental_config=g_prof_config)
	103	+ elif device == 'cuda':
	104	+ from torch.profiler import profile, schedule, tensorboard_trace_handler, ProfilerActivity
	105	+ return profile(
	106	+ activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],
	107	+ schedule=schedule(wait=0, warmup=warm_step, active=active_step, repeat=1),
	108	+ on_trace_ready=tensorboard_trace_handler(profiling_save_path),
	109	+ record_shapes=False,
	110	+ profile_memory=False,
	111	+ )
	112	+ else:
	113	+ print(f"[Warning]: No supported acceleration device (CUDA/NPU) detected, profiler will be disabled")
	114	+ return None
	115	+
	116	+def detect_device_type():
	117	+ try:
	118	+ import torch
	119	+ if torch.cuda.is_available():
	120	+ return "cuda"
	121	+ try:
	122	+ import torch_npu
	123	+ import os
	124	+ os.environ['TORCHINDUCTOR_NPU_BACKEND'] = 'mlir'
	125	+ if torch.npu.is_available():
	126	+ return "npu"
	127	+ except ImportError:
	128	+ pass
	129	+ except ImportError:
	130	+ pass
	131	+ return "cpu"
	132	+
	133	+def device_synchronize():
	134	+ if torch.cuda.is_available():
	135	+ torch.cuda.synchronize()
	136	+ elif torch.npu.is_available():
	137	+ torch.npu.synchronize()
	138	+
	139	+def eval_op_prof(model: nn.Module, device, mod, args):
	140	+ dl = iter(val_loader)
	141	+ x, _ = next(dl)
	142	+ x = x.to(device)
	143	+ pt_path= MODEL_NAME +'_'+ mod + '_result.pt'
	144	+
	145	+ model.eval()
	146	+ with torch.no_grad():
	147	+ res = model(x)
	148	+ torch.save(res, pt_path)
	149	+ device_synchronize()
	150	+
	151	+ prof = None
	152	+ if args.enable_profiler:
	153	+ profiling_save_path = args.profiler_save_path + '/' + MODEL_NAME + '/' + mod
	154	+ prof = get_profile(args.profiler_start_step, args.profiler_end_step, profiling_save_path)
	155	+ prof.start()
	156	+
	157	+ execution_times = []
	158	+ with torch.no_grad():
	159	+ for i in range(args.max_steps):
	160	+ start_time = time.time()
	161	+ model(x)
	162	+ device_synchronize()
	163	+ end_time = time.time()
	164	+ step_time_ms = (end_time - start_time) * 1000
	165	+ print(f"[{mod}]: Step {i}: {step_time_ms:.4f} ms")
	166	+
	167	+ if i >= 10:
	168	+ execution_times.append(step_time_ms)
	169	+
	170	+ if args.enable_profiler:
	171	+ prof.step()
	172	+
	173	+ if args.enable_profiler:
	174	+ prof.stop()
	175	+ if execution_times:
	176	+ avg_ms = sum(execution_times) / len(execution_times)
	177	+ print(f"[{mod}]:: Avg over {len(execution_times)} steps: {avg_ms:.4f} ms")
	178	+
	179	+parser = argparse.ArgumentParser(description=MODEL_NAME + " infernece")
	180	+parser.add_argument("--max_steps", type=int, default=200,
	181	+ help="Total training steps")
	182	+parser.add_argument("--enable_compile", action="store_true",
	183	+ help="Enable torch.compile and Inductor backend")
	184	+parser.add_argument("--enable_profiler", action="store_true",
	185	+ help="Enable profiler for performance analysis")
	186	+parser.add_argument("--profiler_start_step", type=int, default=5,
	187	+ help="Output directory for trained model")
	188	+parser.add_argument("--profiler_end_step", type=int, default=8,
	189	+ help="Output directory for trained model")
	190	+parser.add_argument("--profiler_save_path", type=str, default="./profile",
	191	+ help="Output directory for trained model")
	192	+args = parser.parse_args()
	193	+
	194	+from CTR_Algorithm.pytorch.model.DCNv2 import DeepCrossNetv2
	195	+model = DeepCrossNetv2(feature_fields = fields, embed_dim = 16, layer_num = 2, mlp_dims = (32, 16), dropout = 0.1, cross_method='Matrix')
	196	+
	197	+device_type = detect_device_type()
	198	+device = torch.device(device_type)
	199	+model.to(device)
	200	+mod = 'eager'
	201	+if args.enable_compile:
	202	+ model = torch.compile(model, dynamic=False)
	203	+ mod = 'compile'
	204	+print(f"{MODEL_NAME} inference begin, mode is {mod}")
	205	+eval_op_prof(model, device, mod, args)

Abenchmarks/inference/DIN/din.py+176-0

@@ -0,0 +1,176 @@
	1	+# -- coding: utf-8 --
	2	+import os
	3	+import time
	4	+import argparse
	5	+import sys
	6	+import pandas as pd
	7	+import numpy as np
	8	+import torch
	9	+import torch.nn as nn
	10	+import torch.optim as optim
	11	+import torch.utils.data as Data
	12	+import torch.nn.functional as F
	13	+from sklearn.model_selection import train_test_split
	14	+
	15	+from CTR_Algorithm.pytorch.model import DIN
	16	+from CTR_Algorithm.data.AmazonDataPreprocess import AmazonBookPreprocess
	17	+
	18	+MODEL_NAME = "DIN"
	19	+
	20	+data = pd.read_csv('CTR_Algorithm/data/amazon-books-100k.txt')
	21	+data = AmazonBookPreprocess(data)
	22	+fields = data.max().max()
	23	+
	24	+data_X = data.iloc[:,:-1]
	25	+data_y = data.label.values
	26	+
	27	+tmp_X, test_X, tmp_y, test_y = train_test_split(data_X, data_y, test_size = 0.2, random_state=42, stratify=data_y)
	28	+train_X, val_X, train_y, val_y = train_test_split(tmp_X, tmp_y, test_size = 0.25, random_state=42, stratify=tmp_y)
	29	+
	30	+train_X = torch.from_numpy(train_X.values).long()
	31	+val_X = torch.from_numpy(val_X.values).long()
	32	+test_X = torch.from_numpy(test_X.values).long()
	33	+
	34	+train_y = torch.from_numpy(train_y).long()
	35	+val_y = torch.from_numpy(val_y).long()
	36	+test_y = torch.from_numpy(test_y).long()
	37	+
	38	+train_set = Data.TensorDataset(train_X, train_y)
	39	+val_set = Data.TensorDataset(val_X, val_y)
	40	+train_loader = Data.DataLoader(dataset=train_set, batch_size=32, shuffle=True)
	41	+val_loader = Data.DataLoader(dataset=val_set, batch_size=32, shuffle=False)
	42	+
	43	+def get_profile(profiler_start_step: int, profiler_end_step: int, profiling_save_path: str):
	44	+ warm_step = profiler_start_step
	45	+ active_step = profiler_end_step - warm_step +1
	46	+ print(f"[Profile INFO]: warm_step: {warm_step}, active_step: {active_step}, profiling_save_path: {profiling_save_path}")
	47	+ device = detect_device_type()
	48	+ if device == 'npu':
	49	+ import torch_npu
	50	+ g_prof_config = torch_npu.profiler._ExperimentalConfig(
	51	+ export_type=[
	52	+ torch_npu.profiler.ExportType.Text,
	53	+ torch_npu.profiler.ExportType.Db
	54	+ ],
	55	+ profiler_level=torch_npu.profiler.ProfilerLevel.Level2,
	56	+ msprof_tx=False,
	57	+ aic_metrics=torch_npu.profiler.AiCMetrics.AiCoreNone,
	58	+ l2_cache=False,
	59	+ op_attr=False,
	60	+ data_simplification=False,
	61	+ record_op_args=False,
	62	+ gc_detect_threshold=None)
	63	+
	64	+ return torch_npu.profiler.profile(
	65	+ activities=[
	66	+ torch_npu.profiler.ProfilerActivity.CPU,
	67	+ torch_npu.profiler.ProfilerActivity.NPU],
	68	+ schedule=torch_npu.profiler.schedule(wait=0, warmup=warm_step, active=active_step, repeat=1),
	69	+ on_trace_ready=torch_npu.profiler.tensorboard_trace_handler(profiling_save_path),
	70	+ record_shapes=False,
	71	+ profile_memory=False,
	72	+ with_stack=False,
	73	+ with_modules=False,
	74	+ with_flops=False,
	75	+ experimental_config=g_prof_config)
	76	+ elif device == 'cuda':
	77	+ from torch.profiler import profile, schedule, tensorboard_trace_handler, ProfilerActivity
	78	+ return profile(
	79	+ activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],
	80	+ schedule=schedule(wait=0, warmup=warm_step, active=active_step, repeat=1),
	81	+ on_trace_ready=tensorboard_trace_handler(profiling_save_path),
	82	+ record_shapes=False,
	83	+ profile_memory=False,
	84	+ )
	85	+ else:
	86	+ print(f"[Warning]: No supported acceleration device (CUDA/NPU) detected, profiler will be disabled")
	87	+ return None
	88	+
	89	+def detect_device_type():
	90	+ try:
	91	+ import torch
	92	+ if torch.cuda.is_available():
	93	+ return "cuda"
	94	+ try:
	95	+ import torch_npu
	96	+ import os
	97	+ os.environ['TORCHINDUCTOR_NPU_BACKEND'] = 'mlir'
	98	+ if torch.npu.is_available():
	99	+ return "npu"
	100	+ except ImportError:
	101	+ pass
	102	+ except ImportError:
	103	+ pass
	104	+ return "cpu"
	105	+
	106	+def device_synchronize():
	107	+ if torch.cuda.is_available():
	108	+ torch.cuda.synchronize()
	109	+ elif torch.npu.is_available():
	110	+ torch.npu.synchronize()
	111	+
	112	+def eval_op_prof(model: nn.Module, device, mod, args):
	113	+ dl = iter(val_loader)
	114	+ x, _ = next(dl)
	115	+ x = x.to(device)
	116	+ pt_path= MODEL_NAME +'_'+ mod + '_result.pt'
	117	+
	118	+ model.eval()
	119	+ with torch.no_grad():
	120	+ res = model(x)
	121	+ torch.save(res, pt_path)
	122	+ device_synchronize()
	123	+
	124	+ prof = None
	125	+ if args.enable_profiler:
	126	+ profiling_save_path = args.profiler_save_path + '/' + MODEL_NAME + '/' + mod
	127	+ prof = get_profile(args.profiler_start_step, args.profiler_end_step, profiling_save_path)
	128	+ prof.start()
	129	+
	130	+ execution_times = []
	131	+ with torch.no_grad():
	132	+ for i in range(args.max_steps):
	133	+ start_time = time.time()
	134	+ model(x)
	135	+ device_synchronize()
	136	+ end_time = time.time()
	137	+ step_time_ms = (end_time - start_time) * 1000
	138	+ print(f"[{mod}]: Step {i}: {step_time_ms:.4f} ms")
	139	+
	140	+ if i >= 10:
	141	+ execution_times.append(step_time_ms)
	142	+
	143	+ if args.enable_profiler:
	144	+ prof.step()
	145	+
	146	+ if args.enable_profiler:
	147	+ prof.stop()
	148	+ if execution_times:
	149	+ avg_ms = sum(execution_times) / len(execution_times)
	150	+ print(f"[{mod}]:: Avg over {len(execution_times)} steps: {avg_ms:.4f} ms")
	151	+
	152	+parser = argparse.ArgumentParser(description=MODEL_NAME + " infernece")
	153	+parser.add_argument("--max_steps", type=int, default=200,
	154	+ help="Total training steps")
	155	+parser.add_argument("--enable_compile", action="store_true",
	156	+ help="Enable torch.compile and Inductor backend")
	157	+parser.add_argument("--enable_profiler", action="store_true",
	158	+ help="Enable profiler for performance analysis")
	159	+parser.add_argument("--profiler_start_step", type=int, default=5,
	160	+ help="Output directory for trained model")
	161	+parser.add_argument("--profiler_end_step", type=int, default=8,
	162	+ help="Output directory for trained model")
	163	+parser.add_argument("--profiler_save_path", type=str, default="./profile",
	164	+ help="Output directory for trained model")
	165	+args = parser.parse_args()
	166	+
	167	+model = DIN.DeepInterestNet(feature_dim=fields, embed_dim=8, mlp_dims=[64,32], dropout=float(0.2))
	168	+device_type = detect_device_type()
	169	+device = torch.device(device_type)
	170	+model.to(device)
	171	+mod = 'eager'
	172	+if args.enable_compile:
	173	+ model = torch.compile(model, dynamic=False)
	174	+ mod = 'compile'
	175	+print(f"{MODEL_NAME} inference begin, mode is {mod}")
	176	+eval_op_prof(model, device, mod, args)

Abenchmarks/inference/ETA/eta.py+510-0

@@ -0,0 +1,510 @@
	1	+import argparse
	2	+import glob
	3	+import json
	4	+import os
	5	+import random
	6	+import stat
	7	+import time
	8	+from datetime import datetime
	9	+
	10	+import h5py
	11	+import numpy as np
	12	+
	13	+import torch
	14	+import torch.nn as nn
	15	+import torch.optim as optim
	16	+import torch.nn.functional as F
	17	+from torch.nn.utils.rnn import pad_sequence
	18	+from torch.utils.data import Dataset, DataLoader, ConcatDataset
	19	+from sklearn.metrics import roc_auc_score
	20	+
	21	+torch.manual_seed(2024)
	22	+random.seed(2024)
	23	+
	24	+MODEL_NAME = "ETA"
	25	+STD_DEV = (2 / 512) ** 0.5
	26	+EMBEDDING_FEATURE_NUM = 27
	27	+TARGET_FIELDS = ["206", "207", "216", "210"]
	28	+
	29	+def detect_device_type():
	30	+ try:
	31	+ import torch
	32	+ if torch.cuda.is_available():
	33	+ return "cuda"
	34	+ try:
	35	+ import torch_npu
	36	+ import os
	37	+ os.environ['TORCHINDUCTOR_NPU_BACKEND'] = 'mlir'
	38	+ if torch.npu.is_available():
	39	+ return "npu"
	40	+ except ImportError:
	41	+ pass
	42	+ except ImportError:
	43	+ pass
	44	+ return "cpu"
	45	+
	46	+def json_file_load(json_name: str, json_path: str) -> dict:
	47	+ """
	48	+ Load a JSON file from the specified path.
	49	+ """
	50	+ flags = os.O_RDONLY
	51	+ modes = stat.S_IRUSR \| stat.S_IWUSR \| stat.S_IRGRP \| stat.S_IROTH
	52	+ try:
	53	+ with os.fdopen(os.open(json_path, flags, modes), "r") as fp:
	54	+ json_re = json.load(fp)
	55	+ except FileNotFoundError as e:
	56	+ raise FileNotFoundError(f"{json_name} file not found: {e}") from e
	57	+ except Exception as e:
	58	+ raise RuntimeError(f"Error loading {json_name} file: {e}") from e
	59	+
	60	+ return json_re
	61	+
	62	+def collate_fn(batch):
	63	+ input_dicts = [item[0] for item in batch]
	64	+ target_dicts = [item[1] for item in batch]
	65	+ input_tensors = {}
	66	+ for key in input_dicts[0].keys():
	67	+ tensors = [d[key] for d in input_dicts if key in d and d[key] is not None]
	68	+ if not tensors:
	69	+ continue
	70	+ if tensors[0].dim() == 0:
	71	+ tensors = [t.unsqueeze(0) for t in tensors]
	72	+ input_tensors[key] = pad_sequence(tensors, batch_first=True)
	73	+
	74	+ target_tensors = {}
	75	+ for key in target_dicts[0].keys():
	76	+ tensors = [d[key] for d in target_dicts if key in d and d[key] is not None]
	77	+ if not tensors:
	78	+ continue
	79	+ target_tensors[key] = torch.stack(tensors)
	80	+
	81	+ return input_tensors, target_tensors
	82	+
	83	+def parse_arguments():
	84	+ parser = argparse.ArgumentParser()
	85	+ parser.add_argument("--embedding_size", type=int, default=16, help="Embedding size")
	86	+ parser.add_argument("--attention_dim", type=int, default=64, help="")
	87	+ parser.add_argument("--enable_compile", action="store_true", help="")
	88	+ parser.add_argument("--num_heads", type=int, default=4, help="")
	89	+ parser.add_argument("--short_output_dim", type=int, default=16, help="short attention output dimension")
	90	+ parser.add_argument("--max_seq_len", type=int, default=50, help="")
	91	+ parser.add_argument("--topk", type=int, default=16, help="")
	92	+ parser.add_argument("--deep_layers", type=str, default="512,256,128,64", help="deep layers")
	93	+ parser.add_argument("--hash_bits", type=int, default=32, help="")
	94	+ parser.add_argument("--reuse_hash", type=bool, default=True, help="")
	95	+ parser.add_argument("--batch_size", type=int, default=4096, help="Number of batch size")
	96	+ parser.add_argument("--learning_rate", type=float, default=0.001, help="learning rate")
	97	+ parser.add_argument("--optimizer", type=str, default="Adam",
	98	+ choices=["Adam", "Adagrad", "GD", "Momentum"], help="")
	99	+ parser.add_argument('--early_stop_patience', type=int, default=5, help="")
	100	+ parser.add_argument("--data_dir", type=str, default="./aliccp/aliccp_out/", help="data dir")
	101	+ parser.add_argument("--dt_dir", type=str, default='', help="data dt partition")
	102	+ parser.add_argument("--model_dir", type=str, default=f"./", help="code check point dir")
	103	+ parser.add_argument("--clear_existing_model", action="store_true", help="")
	104	+ parser.add_argument("--task_type", type=str, default="train",
	105	+ choices=["train", "eval", "predict"], help="task type")
	106	+ parser.add_argument("--max_steps", type=int, default=100,
	107	+ help="Total training steps (优先级高于num_epochs，-1表示使用num_epochs控制)")
	108	+ parser.add_argument('--epoch_num', type=int, default=1, help="Number of epochs")
	109	+ parser.add_argument('--train_batch_num', type=int, default=2000, help="Number of train batchs")
	110	+ parser.add_argument('--test_batch_num', type=int, default=10, help="Number of test batchs")
	111	+ parser.add_argument("--enable_profiler", action="store_true",
	112	+ help="Enable profiler for performance analysis")
	113	+ parser.add_argument("--profiler_start_step", type=int, default=5,
	114	+ help="Output directory for trained model")
	115	+ parser.add_argument("--profiler_end_step", type=int, default=8,
	116	+ help="Output directory for trained model")
	117	+ parser.add_argument("--profiler_save_path", type=str, default="./profile",
	118	+ help="Output directory for trained model")
	119	+ return parser.parse_args()
	120	+
	121	+
	122	+class HDF5Dataset(Dataset):
	123	+ def __init__(self, hdf5_path):
	124	+ self.hdf5_path = hdf5_path
	125	+ self._load_hdf5()
	126	+
	127	+ def _load_hdf5(self):
	128	+ with h5py.File(self.hdf5_path, 'r') as f:
	129	+ self.input_sample = {}
	130	+ self.target_sample = {}
	131	+ y = np.array(f["y"])
	132	+ z = np.array(f["z"])
	133	+ fields = [key for key in f.keys() if key not in ["y", "z"]]
	134	+ self.target_sample.update({"y": torch.tensor(y, dtype=torch.float32)})
	135	+ self.target_sample.update({"z": torch.tensor(z, dtype=torch.float32)})
	136	+ for multi_field in fields:
	137	+ self.input_sample.update({multi_field: torch.tensor(np.array(f[multi_field]), dtype=torch.int64)})
	138	+ self._length = len(y)
	139	+
	140	+ def __len__(self):
	141	+ return self._length
	142	+
	143	+ def __getitem__(self, idx):
	144	+ input_dict = {k: v[idx] for k, v in self.input_sample.items()}
	145	+ target_dict = {k: v[idx] for k, v in self.target_sample.items()}
	146	+ return input_dict, target_dict
	147	+
	148	+
	149	+class TorchDataSet(ConcatDataset):
	150	+ def __init__(self, files):
	151	+ datasets = [HDF5Dataset(fp) for fp in files]
	152	+ super().__init__(datasets)
	153	+
	154	+
	155	+class ETA(nn.Module):
	156	+ def __init__(self, spec, params):
	157	+ super(ETA, self).__init__()
	158	+ self.spec = spec
	159	+ self.params = params
	160	+
	161	+ # embedding layers
	162	+ self.emb_weights = nn.ModuleDict()
	163	+ for key, vocab_len in spec["vocab_length"].items():
	164	+ self.emb_weights[key] = nn.Embedding(vocab_len + 1, params.embedding_size)
	165	+ nn.init.normal_(self.emb_weights[key].weight, std=STD_DEV)
	166	+
	167	+ self.hash_weights = nn.Parameter(
	168	+ torch.randn(params.embedding_size, params.hash_bits),
	169	+ requires_grad=False
	170	+ )
	171	+
	172	+ # attention paramters
	173	+ self.short_attentions = nn.ModuleList()
	174	+ self.long_attentions = nn.ModuleList()
	175	+ for _ in range(len(TARGET_FIELDS)): # 4 target fields
	176	+ self.short_attentions.append(ShortAttention(params))
	177	+ self.long_attentions.append(LongAttention(self.hash_weights, params))
	178	+
	179	+ # mlp layer
	180	+ self.mlp = nn.Sequential()
	181	+ deep_layers = list(map(int, params.deep_layers.strip().split(",")))
	182	+ input_dim = EMBEDDING_FEATURE_NUM * params.embedding_size
	183	+ for i, dim in enumerate(deep_layers):
	184	+ self.mlp.add_module(f'mlp{i}', nn.Linear(input_dim, dim))
	185	+ self.mlp.add_module(f'relu{i}', nn.ReLU())
	186	+ input_dim = dim
	187	+
	188	+ # output layer
	189	+ self.output = nn.Linear(input_dim, 1)
	190	+
	191	+ def embedding_lookup_sparse(self, params: nn.Embedding, ids: torch.Tensor, combiner: str):
	192	+ mask = (ids >= 0).float().unsqueeze(-1)
	193	+ ids = ids.clone()
	194	+ ids[ids == -1] = 0
	195	+ embedding = params(ids) * mask
	196	+ if combiner == "sum":
	197	+ return embedding.sum(dim=1)
	198	+ elif combiner == "mean":
	199	+ return embedding.sum(dim=1) / mask.sum(dim=1).clamp(min=1e-7)
	200	+ else:
	201	+ raise ValueError("combiner only supoort 'sum', 'mean'")
	202	+
	203	+ def forward(self, features):
	204	+ # embedding
	205	+ embeddings = {}
	206	+ masks = {}
	207	+ for key in self.spec["one_hot_fields"]:
	208	+ embeddings[key] = self.emb_weights[key](features[key]).unsqueeze(1)
	209	+ for key in self.spec["multi_hot_fields"]:
	210	+ feat = features[key]
	211	+ masks[key] = (feat >= 0).bool().unsqueeze(1)
	212	+ feat = feat.clone()
	213	+ feat[feat == -1] = 0
	214	+ embeddings[key] = self.emb_weights[key](feat)
	215	+ for key in self.spec["special_fields"]:
	216	+ embeddings[key] = self.embedding_lookup_sparse(
	217	+ self.emb_weights[key], features[key], combiner="sum"
	218	+ ).unsqueeze(1)
	219	+
	220	+ def long_emb_cat(field_name):
	221	+ dense_embedding = embeddings.get(field_name)
	222	+ dense_mask = masks.get(field_name)
	223	+ if dense_embedding is None or dense_mask is None:
	224	+ raise ValueError(f"Field {field_name} not found in embeddings or masks")
	225	+ padded_embedding = F.pad(dense_embedding, (0, 0, 0, self.params.max_seq_len), "constant", 0)
	226	+ padded_mask = F.pad(dense_mask, (0, self.params.max_seq_len), "constant", False)
	227	+ return (
	228	+ padded_embedding[:, :self.params.topk, :],
	229	+ padded_embedding[:, :self.params.max_seq_len, :],
	230	+ padded_mask[:, :, :self.params.topk],
	231	+ padded_mask[:, :, :self.params.max_seq_len]
	232	+ )
	233	+ emb_cats = [long_emb_cat(field) for field in self.spec["multi_hot_fields"]]
	234	+ target_fields = TARGET_FIELDS
	235	+
	236	+ short_attns = []
	237	+ long_attns = []
	238	+
	239	+ for i, (emb_cat, target) in enumerate(zip(emb_cats, target_fields)):
	240	+ emb_target = embeddings[target]
	241	+
	242	+ emb_short = emb_cat[0]
	243	+ mask_short = emb_cat[2]
	244	+ short_attns.append(self.short_attentions[i](emb_target, emb_short, mask_short))
	245	+
	246	+ emb_long = emb_cat[1]
	247	+ mask_long = emb_cat[3]
	248	+ long_attns.append(self.long_attentions[i](emb_target, emb_long, mask_long))
	249	+
	250	+ # concat all embeddings
	251	+ all_embs = []
	252	+ for field in self.spec["one_hot_fields"] + self.spec["special_fields"]:
	253	+ if embeddings[field].dim() > 3:
	254	+ all_embs.append(embeddings[field].squeeze(1))
	255	+ else:
	256	+ all_embs.append(embeddings[field])
	257	+ all_embs += short_attns + long_attns
	258	+ all_embs = torch.cat(all_embs, dim=1)
	259	+ if all_embs.dim() >= 2:
	260	+ concat_emb = all_embs.squeeze(1)
	261	+ else:
	262	+ concat_emb = all_embs
	263	+
	264	+ concat_emb = concat_emb.reshape(concat_emb.size(0), -1)
	265	+ # mlp and outputs
	266	+ mlp_out = self.mlp(concat_emb)
	267	+ logits = self.output(mlp_out).squeeze()
	268	+ pred = torch.sigmoid(logits)
	269	+ return pred, logits
	270	+
	271	+ def build_loss(self, pred, labels, click_weight=0.14, epsilon=1e-7):
	272	+ if pred.shape != labels.shape:
	273	+ raise ValueError(f"pred and labels must be the same shape. "
	274	+ f"pred shape: {pred.shape}, labels shape: {labels.shape}")
	275	+ pred = torch.clamp(pred, min=epsilon, max=1 - epsilon)
	276	+ loss = - (1 - click_weight) / click_weight * labels * torch.log(pred) - (1 - labels) * torch.log(1 - pred)
	277	+ return loss.mean()
	278	+
	279	+ def build_optimizer(self):
	280	+ if self.params.optimizer == "Adam":
	281	+ optimizer = optim.Adam(
	282	+ params=self.parameters(),
	283	+ lr=self.params.learning_rate,
	284	+ betas=[0.9, 0.999], eps=1e-8
	285	+ )
	286	+ elif self.params.optimizer == "Adagrad":
	287	+ optimizer = optim.Adagrad(
	288	+ params=self.parameters(),
	289	+ lr=self.params.learning_rate,
	290	+ initial_accumulator_value=1e-6
	291	+ )
	292	+ elif self.params.optimizer == "Momentum":
	293	+ optimizer = optim.SGD(
	294	+ params=self.parameters(),
	295	+ lr=self.params.learning_rate,
	296	+ momentum=0.95
	297	+ )
	298	+ elif self.params.optimizer == "SGD":
	299	+ optimizer = optim.SGD(
	300	+ params=self.parameters(),
	301	+ lr=self.params.learning_rate,
	302	+ )
	303	+ else:
	304	+ raise ValueError("Unsupported optimizer type: {}".format(args.optimizer))
	305	+ return optimizer
	306	+
	307	+
	308	+class ShortAttention(nn.Module):
	309	+ def __init__(self, params):
	310	+ super().__init__()
	311	+ self.attention_dim = params.attention_dim
	312	+ self.num_heads = params.num_heads
	313	+ self.key_dim = self.attention_dim // self.num_heads
	314	+
	315	+ self.q_fc = nn.Linear(params.embedding_size, self.attention_dim)
	316	+ self.k_fc = nn.Linear(params.embedding_size, self.attention_dim)
	317	+ self.v_fc = nn.Linear(params.embedding_size, self.attention_dim)
	318	+ self.o_fc = nn.Linear(self.attention_dim, params.short_output_dim)
	319	+
	320	+ def forward(self, target_input, seq_input, mask):
	321	+ # target_input:[B, 1, E], seq_input:[B, S, E], mask:[B, 1, S] B = target_input.size(0)
	322	+ # logger.info(f"tgt shape: {target_input.shape}, seq shape: {seq_input.shape}, mask shape: {mask.shape}")
	323	+ b, tgt_len = target_input.shape[:2]
	324	+ b, seq_len = seq_input.shape[:2]
	325	+ query = self.q_fc(target_input) # [B, 1, A]
	326	+ key = self.k_fc(seq_input) # [B, S, A]
	327	+ value = self.v_fc(seq_input) # [B, S, A]
	328	+
	329	+ # split heads
	330	+ query = query.view(b, tgt_len, self.num_heads, self.key_dim).permute(0, 2, 1, 3) # [B, Heads, 1, key_dim]
	331	+ key = key.view(b, seq_len, self.num_heads, self.key_dim).permute(0, 2, 3, 1) # [B, Heads, key_dim, S]
	332	+ value = value.view(b, seq_len, self.num_heads, self.key_dim).permute(0, 2, 1, 3) # [B, Heads, S, key_dim]
	333	+
	334	+ # scaled dot-product attention
	335	+ scores = torch.matmul(query, key) / (self.key_dim ** 0.5) # [B, Heads, 1, S]
	336	+ scores = scores.masked_fill(~mask.unsqueeze(1), float('-inf'))
	337	+
	338	+ attn = F.softmax(scores, dim=-1)
	339	+
	340	+ output = torch.matmul(attn, value) # [B, Heads, 1, key_dim]
	341	+ output = output.permute(0, 2, 1, 3)
	342	+ b, lens, heads, dims = output.shape
	343	+ output = output.reshape(b * lens, 1, heads * dims)
	344	+ output = self.o_fc(output)
	345	+ return output
	346	+
	347	+
	348	+class LongAttention(nn.Module):
	349	+ def __init__(self, hash_weights, params):
	350	+ super().__init__()
	351	+ self.hash_weights = hash_weights
	352	+ self.reuse_hash = params.reuse_hash
	353	+ self.topk = params.topk
	354	+ self.short_attn = ShortAttention(params)
	355	+
	356	+ def lsh_hash(self, vecs, hash_weights):
	357	+ rotated_vecs = torch.matmul(vecs, hash_weights)
	358	+ return (rotated_vecs > 0).float()
	359	+
	360	+ def forward(self, target_input, seq_input, mask):
	361	+ # target_input:[B, 1, E], seq_input:[B, S, E], mask:[B, 1, S]
	362	+ b, s, e = seq_input.shape
	363	+ hash_weights = self.hash_weights if self.reuse_hash else nn.Parameter(
	364	+ torch.randn(e, self.hash_weights.size(-1)),
	365	+ requires_grad=False
	366	+ ) # [E, hash_bits]
	367	+
	368	+ target_hash = self.lsh_hash(target_input.squeeze(), hash_weights) # [B, hash_bits]
	369	+ seq_hash = self.lsh_hash(seq_input, hash_weights) # [B, S, hash_bits]
	370	+ # calculate similarity(hamming distance)
	371	+ hash_sim = - torch.sum(torch.abs(seq_hash - target_hash.unsqueeze(1)), dim=-1) # [B, S]
	372	+ hash_sim = hash_sim.masked_fill(~mask.squeeze(1), float('-inf'))
	373	+
	374	+ # select topk
	375	+ _, topk_idx = torch.topk(hash_sim, self.topk, dim=-1) # [B, topk]
	376	+
	377	+ topk_seq = torch.gather(
	378	+ seq_input,
	379	+ 1,
	380	+ topk_idx.unsqueeze(-1).expand(-1, -1, e)
	381	+ )
	382	+
	383	+ topk_idx = topk_idx.unsqueeze(1)
	384	+ topk_mask = torch.gather(mask, 2, topk_idx)
	385	+
	386	+ # apply short attention to topk
	387	+ output = self.short_attn(target_input, topk_seq, topk_mask)
	388	+ return output
	389	+
	390	+def get_profile(profiler_start_step: int, profiler_end_step: int, profiling_save_path: str):
	391	+ warm_step = profiler_start_step
	392	+ active_step = profiler_end_step - warm_step +1
	393	+ print(f"[Profile INFO]: warm_step: {warm_step}, active_step: {active_step}, profiling_save_path: {profiling_save_path}")
	394	+ device = detect_device_type()
	395	+ if device == 'npu':
	396	+ import torch_npu
	397	+ g_prof_config = torch_npu.profiler._ExperimentalConfig(
	398	+ export_type=[
	399	+ torch_npu.profiler.ExportType.Text,
	400	+ torch_npu.profiler.ExportType.Db
	401	+ ],
	402	+ profiler_level=torch_npu.profiler.ProfilerLevel.Level2,
	403	+ msprof_tx=False,
	404	+ aic_metrics=torch_npu.profiler.AiCMetrics.AiCoreNone,
	405	+ l2_cache=False,
	406	+ op_attr=False,
	407	+ data_simplification=False,
	408	+ record_op_args=False,
	409	+ gc_detect_threshold=None)
	410	+
	411	+ return torch_npu.profiler.profile(
	412	+ activities=[
	413	+ torch_npu.profiler.ProfilerActivity.CPU,
	414	+ torch_npu.profiler.ProfilerActivity.NPU],
	415	+ schedule=torch_npu.profiler.schedule(wait=0, warmup=warm_step, active=active_step, repeat=1),
	416	+ on_trace_ready=torch_npu.profiler.tensorboard_trace_handler(profiling_save_path),
	417	+ record_shapes=False,
	418	+ profile_memory=False,
	419	+ with_stack=False,
	420	+ with_modules=False,
	421	+ with_flops=False,
	422	+ experimental_config=g_prof_config)
	423	+ elif device == 'cuda':
	424	+ from torch.profiler import profile, schedule, tensorboard_trace_handler, ProfilerActivity
	425	+ return profile(
	426	+ activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],
	427	+ schedule=schedule(wait=0, warmup=warm_step, active=active_step, repeat=1),
	428	+ on_trace_ready=tensorboard_trace_handler(profiling_save_path),
	429	+ record_shapes=False,
	430	+ profile_memory=False,
	431	+ )
	432	+ else:
	433	+ print(f"[Warning]: No supported acceleration device (CUDA/NPU) detected, profiler will be disabled")
	434	+ return None
	435	+
	436	+
	437	+def device_synchronize():
	438	+ if torch.cuda.is_available():
	439	+ torch.cuda.synchronize()
	440	+ elif torch.npu.is_available():
	441	+ torch.npu.synchronize()
	442	+
	443	+
	444	+def evaluate_op(model: ETA, dataloader, device, args):
	445	+ mod = 'compile' if args.enable_compile else 'eager'
	446	+ dl = iter(dataloader)
	447	+ features, _ = next(dl)
	448	+ features = {k: v.to(device) for k, v in features.items()}
	449	+
	450	+ for key, value in features.items():
	451	+ if isinstance(value, torch.Tensor) and value.dtype == torch.int64:
	452	+ features[key] = value.to(torch.int32)
	453	+
	454	+ if mod == 'compile':
	455	+ model = torch.compile(model, dynamic=False)
	456	+ model.eval()
	457	+ with torch.no_grad():
	458	+ model(features)
	459	+ device_synchronize()
	460	+
	461	+ prof = None
	462	+ if args.enable_profiler:
	463	+ profiling_save_path = args.profiler_save_path + '/' + MODEL_NAME + '/' + mod
	464	+ prof = get_profile(args.profiler_start_step, args.profiler_end_step, profiling_save_path)
	465	+ prof.start()
	466	+
	467	+ exec_times = []
	468	+ with torch.no_grad():
	469	+ for i in range(args.max_steps):
	470	+ e2etime1 = time.perf_counter()
	471	+ model(features)
	472	+ device_synchronize()
	473	+ e2etime2 = time.perf_counter()
	474	+ elapsed_ms = (e2etime2 - e2etime1) * 1000
	475	+ print(f"iterations {i} / {args.max_steps}: [{mod}] e2e time: {elapsed_ms} ms.")
	476	+
	477	+ if i >= 10:
	478	+ exec_times.append(elapsed_ms)
	479	+
	480	+ if args.enable_profiler:
	481	+ prof.step()
	482	+ if args.enable_profiler:
	483	+ prof.stop()
	484	+ if exec_times:
	485	+ avg_time = sum(exec_times) / len(exec_times)
	486	+ print("Step time consumption statistics (excluding the first 10 steps)")
	487	+ print(f"\n>>> [{mod}] avg e2e: {avg_time:.3f} ms <<<")
	488	+
	489	+def main(args):
	490	+ test_files = glob.glob("%stest/data_test.csv.hd5.*" % args.data_dir)
	491	+ spec = json_file_load("spec", os.path.join(args.data_dir, "spec.json"))
	492	+ model = ETA(spec, args)
	493	+ device_type = detect_device_type()
	494	+ print(f"Using device type {device_type}")
	495	+ device = torch.device(device_type)
	496	+ model.to(device)
	497	+ test_dataset = TorchDataSet(test_files[1:2])
	498	+ test_loader = DataLoader(dataset=test_dataset,
	499	+ batch_size=args.batch_size,
	500	+ shuffle=True,
	501	+ collate_fn=collate_fn,
	502	+ prefetch_factor=100,
	503	+ num_workers=10)
	504	+ evaluate_op(model, test_loader, device, args)
	505	+
	506	+
	507	+if __name__ == "__main__":
	508	+ print(f"{MODEL_NAME} inference begin !!!")
	509	+ args = parse_arguments()
	510	+ main(args)

Abenchmarks/inference/MMOE/mmoe.py+577-0

@@ -0,0 +1,577 @@
	1	+import os
	2	+import stat
	3	+import glob
	4	+import json
	5	+import random
	6	+import shutil
	7	+import logging
	8	+import time
	9	+from datetime import datetime
	10	+import argparse
	11	+
	12	+import pytz
	13	+import h5py
	14	+import numpy as np
	15	+import torch
	16	+import torch.nn as nn
	17	+import torch.optim as optim
	18	+from torch.nn.utils.rnn import pad_sequence
	19	+from torch.utils.data import DataLoader, Dataset, ConcatDataset
	20	+from sklearn.metrics import roc_auc_score
	21	+
	22	+torch.manual_seed(2024)
	23	+random.seed(2024)
	24	+
	25	+MODEL_NAME = "MMOE"
	26	+EMBEDDING_FEATURE_NUM = 23
	27	+STD_DEV = (2 / 512) ** 0.5
	28	+
	29	+def detect_device_type():
	30	+ try:
	31	+ import torch
	32	+ if torch.cuda.is_available():
	33	+ return "cuda"
	34	+ try:
	35	+ import torch_npu
	36	+ import os
	37	+ os.environ['TORCHINDUCTOR_NPU_BACKEND'] = 'mlir'
	38	+ if torch.npu.is_available():
	39	+ return "npu"
	40	+ except ImportError:
	41	+ pass
	42	+ except ImportError:
	43	+ pass
	44	+ return "cpu"
	45	+
	46	+def parse_arguments():
	47	+ parser = argparse.ArgumentParser(description='PyTorch Example with Command Line Arguments')
	48	+ parser.add_argument('--embedding_size', type=int, default=16, help="Embedding size")
	49	+ parser.add_argument('--batch_size', type=int, default=4096, help="Batch size for training")
	50	+ parser.add_argument('--learning_rate', type=float, default=0.001, help="Learning rate")
	51	+ parser.add_argument('--optimizer', type=str, default="Adam", choices=["Adam", "Adagrad", "GD", "Momentum"],
	52	+ help="Optimizer type")
	53	+ parser.add_argument('--expert_layers', type=str, default="512,256", help="Expert layers")
	54	+ parser.add_argument('--tower_layers', type=str, default="128,64", help="tower layers")
	55	+ parser.add_argument('--ctr_task_wgt', type=float, default=0.5, help="loss weight of ctr task")
	56	+ parser.add_argument('--data_dir', type=str, default="./aliccp/aliccp_out/", help="Data directory")
	57	+ parser.add_argument('--dt_dir', type=str, default="", help="Data dt partition")
	58	+ parser.add_argument('--model_dir', type=str, default=f"./",
	59	+ help="Model checkpoint directory")
	60	+ parser.add_argument('--servable_model_dir', type=str, default=f"./",
	61	+ help="Export servable model for pytorch Serving")
	62	+ parser.add_argument('--clear_existing_model', action="store_true", help="Clear existing model or not")
	63	+ parser.add_argument('--max_seq_len', type=int, default=50, help="Max length of sequence")
	64	+ parser.add_argument('--task_num', type=int, default=2, help="Task number")
	65	+ parser.add_argument('--experts_num', type=int, default=8, help="Number of experts")
	66	+ parser.add_argument('--log_level', type=str, default="DEBUG",
	67	+ choices=["DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL"],
	68	+ help="Log level")
	69	+ parser.add_argument('--epoch_num', type=int, default=10, help="Number of epochs")
	70	+ parser.add_argument('--train_batch_num', type=int, default=2000, help="Number of train batchs")
	71	+ parser.add_argument('--eval_batch_num', type=int, default=20, help="Number of eval batchs")
	72	+ parser.add_argument('--test_batch_num', type=int, default=20, help="Number of test batchs")
	73	+ parser.add_argument("--enable_compile", action="store_true", help="")
	74	+ parser.add_argument("--max_steps", type=int, default=100, help="Total training steps")
	75	+ parser.add_argument("--enable_profiler", action="store_true",
	76	+ help="Enable profiler for performance analysis")
	77	+ parser.add_argument("--profiler_start_step", type=int, default=5,
	78	+ help="Output directory for trained model")
	79	+ parser.add_argument("--profiler_end_step", type=int, default=8,
	80	+ help="Output directory for trained model")
	81	+ parser.add_argument("--profiler_save_path", type=str, default="./profile",
	82	+ help="Output directory for trained model")
	83	+ return parser.parse_args()
	84	+
	85	+
	86	+def json_file_load(json_name: str, json_path: str) -> dict:
	87	+ """
	88	+ Load a JSON file from the specified path.
	89	+ """
	90	+ flags = os.O_RDONLY
	91	+ modes = stat.S_IRUSR \| stat.S_IWUSR \| stat.S_IRGRP \| stat.S_IROTH
	92	+ try:
	93	+ with os.fdopen(os.open(json_path, flags, modes), "r") as fp:
	94	+ json_re = json.load(fp)
	95	+ except FileNotFoundError as e:
	96	+ raise FileNotFoundError(f"{json_name} file not found: {e}") from e
	97	+ except Exception as e:
	98	+ raise RuntimeError(f"Error loading {json_name} file: {e}") from e
	99	+
	100	+ return json_re
	101	+
	102	+
	103	+def pre_deal_dataset(dataset):
	104	+ num_pos = 0
	105	+ num_neg = 0
	106	+ num_pos_y = 0
	107	+ num_neg_y = 0
	108	+ dataset_index = 0
	109	+ for sub_ds in dataset.datasets:
	110	+ z = sub_ds.target_sample["z"]
	111	+ num_pos += (z == 1).sum().item()
	112	+ num_neg += (z == 0).sum().item()
	113	+
	114	+ y = sub_ds.target_sample["y"]
	115	+ num_pos_y += (y == 1).sum().item()
	116	+ num_neg_y += (y == 0).sum().item()
	117	+
	118	+ # logger.info("find pos index %s", dataset_index)
	119	+ dataset_index += 1
	120	+ return np.log(num_neg_y / num_pos_y), np.log(num_neg / num_pos)
	121	+
	122	+
	123	+def collate_fn(batch):
	124	+ input_dicts = [item[0] for item in batch]
	125	+ target_dicts = [item[1] for item in batch]
	126	+ input_tensors = {}
	127	+
	128	+ for key in input_dicts[0].keys():
	129	+ tensors = [d[key] for d in input_dicts if key in d and d[key] is not None]
	130	+ if not tensors:
	131	+ continue
	132	+
	133	+ if tensors[0].dim() == 0:
	134	+ tensors = [t.unsqueeze(0) for t in tensors]
	135	+ input_tensors[key] = pad_sequence(tensors, batch_first=True)
	136	+ target_tensors = {}
	137	+ for key in target_dicts[0].keys():
	138	+ tensors = [d[key] for d in target_dicts if key in d and d[key] is not None]
	139	+ if not tensors:
	140	+ continue
	141	+ target_tensors[key] = torch.stack(tensors)
	142	+ return input_tensors, target_tensors
	143	+
	144	+
	145	+class HDF5Dataset(Dataset):
	146	+ def __init__(self, hdf5_path):
	147	+ self.hdf5_path = hdf5_path
	148	+ self._load_hdf5()
	149	+
	150	+ def _load_hdf5(self):
	151	+ with h5py.File(self.hdf5_path, 'r') as f:
	152	+ self.input_sample = {}
	153	+ self.target_sample = {}
	154	+ y = np.array(f["y"])
	155	+ z = np.array(f["z"])
	156	+ fields = [key for key in f.keys() if key not in ["y", "z"]]
	157	+ self.target_sample.update({"y": torch.tensor(y, dtype=torch.float32)})
	158	+ self.target_sample.update({"z": torch.tensor(z, dtype=torch.float32)})
	159	+ for multi_field in fields:
	160	+ self.input_sample.update({multi_field: torch.tensor(np.array(f[multi_field]), dtype=torch.int64)})
	161	+ self._length = len(y)
	162	+ self.positive_indices = np.where(y == 1)[0]
	163	+ self.negative_indices = np.where(y == 0)[0]
	164	+ # logger.info(f"load file {self.hdf5_path} finished")
	165	+
	166	+ def __len__(self):
	167	+ return self._length
	168	+
	169	+ def __getitem__(self, idx):
	170	+ input_dict = {k: v[idx] for k, v in self.input_sample.items()}
	171	+ target_dict = {k: v[idx] for k, v in self.target_sample.items()}
	172	+ return input_dict, target_dict
	173	+
	174	+
	175	+class TorchDataSet(ConcatDataset):
	176	+ def __init__(self, files):
	177	+ print(files)
	178	+ datasets = [HDF5Dataset(fp) for fp in files]
	179	+ super().__init__(datasets)
	180	+
	181	+
	182	+class TorchMmoeModel(nn.Module):
	183	+ def __init__(self, params) -> None:
	184	+ super().__init__()
	185	+ self.params = params
	186	+ self.spec = json_file_load("spec", os.path.join(self.params.data_dir, "spec.json"))
	187	+ self.embedding_layers = self.build_embedding_layers()
	188	+ self.experts = self.build_experts()
	189	+ self.gates = self.build_gate_networks()
	190	+ self.task_output_layers = self.build_task_output_layers()
	191	+
	192	+ self.tower_names = ['ctr', 'cvr']
	193	+ self.towers = nn.ModuleDict()
	194	+ self.towers_output_layers = nn.ModuleDict()
	195	+ tower_units = list(map(int, self.params.tower_layers.strip().split(',')))
	196	+ input_dim = self.params.experts_num * list(map(int, self.params.expert_layers.strip().split(',')))[-1]
	197	+ for name in self.tower_names:
	198	+ tower_layers = []
	199	+ in_dim = input_dim
	200	+ for out_dim in tower_units:
	201	+ tower_layers.append(nn.Linear(in_dim, out_dim))
	202	+ tower_layers.append(nn.BatchNorm1d(out_dim))
	203	+ tower_layers.append(nn.ReLU())
	204	+ tower_layers.append(nn.Dropout(p=0.2))
	205	+ in_dim = out_dim
	206	+ self.towers[name] = nn.Sequential(*tower_layers)
	207	+ self.towers_output_layers[name] = nn.Linear(in_dim, 1)
	208	+
	209	+ def forward(self, features: dict):
	210	+ # Build the embedding layer
	211	+ x_deep = self.get_embedding(features)
	212	+ experts_out = [expert(x_deep) for expert in self.experts]
	213	+ experts_out = torch.stack(experts_out, dim=1)
	214	+
	215	+ gate_outs = [gate(x_deep) for gate in self.gates]
	216	+
	217	+ task_outputs = []
	218	+ for gate_network in gate_outs:
	219	+ gate_network = gate_network.unsqueeze(-1)
	220	+ task_out = torch.multiply(experts_out, gate_network)
	221	+ task_out_shape = list(task_out.shape)
	222	+ task_outputs.append(torch.reshape(task_out, shape=[-1, task_out_shape[1] * task_out_shape[2]]))
	223	+ return task_outputs
	224	+
	225	+ def build_embedding_layers(self):
	226	+ embeddings = nn.ModuleDict()
	227	+ for key, vocab_len in self.spec["vocab_length"].items():
	228	+ embeddings[key] = nn.Embedding(vocab_len + 1, self.params.embedding_size)
	229	+ nn.init.normal_(embeddings[key].weight, std=STD_DEV)
	230	+ return embeddings
	231	+
	232	+ def build_experts(self):
	233	+
	234	+ experts = nn.ModuleList()
	235	+ expert_units = list(map(int, self.params.expert_layers.strip().split(',')))
	236	+ input_dim = self.params.embedding_size * EMBEDDING_FEATURE_NUM
	237	+
	238	+ for _ in range(self.params.experts_num):
	239	+ expert_layers = []
	240	+ in_features = input_dim
	241	+
	242	+ for out_features in expert_units:
	243	+ expert_layers.append(nn.Linear(in_features, out_features))
	244	+ expert_layers.append(nn.BatchNorm1d(out_features))
	245	+ expert_layers.append(nn.ReLU())
	246	+ expert_layers.append(nn.Dropout(0.2))
	247	+ in_features = out_features
	248	+ experts.append(nn.Sequential(*expert_layers))
	249	+ return experts
	250	+
	251	+ def build_gate_networks(self):
	252	+ gates = nn.ModuleList()
	253	+ input_dim = self.params.embedding_size * EMBEDDING_FEATURE_NUM
	254	+
	255	+ for _ in range(self.params.task_num):
	256	+ gate = nn.Sequential(
	257	+ nn.Linear(input_dim, self.params.experts_num),
	258	+ nn.Softmax(dim=1)
	259	+ )
	260	+ gates.append(gate)
	261	+ return gates
	262	+
	263	+ def build_task_output_layers(self):
	264	+ task_output_layers = nn.ModuleList()
	265	+ expert_units = list(map(int, self.params.expert_layers.strip().split(',')))
	266	+ input_dim = self.params.experts_num * expert_units[-1]
	267	+ tower_units = list(map(int, self.params.tower_layers.strip().split(',')))
	268	+ for _ in range(self.params.task_num):
	269	+ tower = []
	270	+ in_features = input_dim
	271	+ for out_features in tower_units:
	272	+ tower.append(nn.Linear(in_features, out_features))
	273	+ tower.append(nn.ReLU())
	274	+ in_features = out_features
	275	+ task_output_layers.append(nn.Sequential(*tower))
	276	+ return task_output_layers
	277	+
	278	+ def embedding_lookup_sparse_fake(self, key,
	279	+ ids: torch.Tensor,
	280	+ combiner: str = None,
	281	+ name: str = None) -> torch.Tensor:
	282	+ dense_mask = torch.unsqueeze(torch.where(ids >= 0,
	283	+ torch.ones_like(ids, dtype=torch.float32),
	284	+ torch.zeros_like(ids)
	285	+ ),
	286	+ dim=-1
	287	+ )
	288	+
	289	+ # Replace invalid IDs (-1) with zeros
	290	+ ids = torch.where(ids == -1, torch.zeros_like(ids), ids)
	291	+ embedding_layer = self.embedding_layers[key]
	292	+ embedding_output = embedding_layer(ids)
	293	+ embedding = embedding_output * dense_mask
	294	+ summed_embedding = torch.sum(embedding, axis=1)
	295	+ if combiner == "sum":
	296	+ return summed_embedding
	297	+ elif combiner == "mean":
	298	+ return summed_embedding / torch.sum(dense_mask, axis=1)
	299	+ else:
	300	+ raise ValueError("combiner only supoort 'sum', 'mean'")
	301	+
	302	+ def get_embedding(self, features: dict) -> torch.Tensor:
	303	+ """
	304	+ Build the embedding layer for the model.
	305	+
	306	+ Args:
	307	+ features (dict): The input features.
	308	+ Returns:
	309	+ torch.Tensor: The concatenated and reshaped embedding tensor.
	310	+ """
	311	+ on_hot_field_lst = self.spec.get("one_hot_fields")
	312	+ other_field_lst = self.spec.get("multi_hot_fields") + self.spec.get("special_fields")
	313	+
	314	+ embeddings = {}
	315	+ for key in on_hot_field_lst:
	316	+ embedding_layer = self.embedding_layers[key]
	317	+ embeddings[key] = embedding_layer(features[key])
	318	+ embeddings[key] = torch.reshape(embeddings[key], [-1, 1, self.params.embedding_size])
	319	+ for key in other_field_lst:
	320	+ embeddings[key] = torch.unsqueeze(
	321	+ self.embedding_lookup_sparse_fake(key=key, ids=features[key], combiner="sum"),
	322	+ dim=1
	323	+ )
	324	+ embedding = torch.concat(
	325	+ [embeddings.get(field_name) for field_name in self.spec.get("one_hot_fields")] +
	326	+ [embeddings.get(field_name) for field_name in self.spec.get("multi_hot_fields")] +
	327	+ [embeddings.get(field_name) for field_name in self.spec.get("special_fields")],
	328	+ dim=2,
	329	+ )
	330	+
	331	+ return torch.reshape(embedding, [-1, EMBEDDING_FEATURE_NUM * self.params.embedding_size])
	332	+
	333	+ def build_tower(self, tower_input: torch.Tensor, name: str) -> torch.Tensor:
	334	+ """
	335	+ Build the tower network for a specific task.
	336	+
	337	+ Args:
	338	+ tower_input (torch.Tensor): The input tensor for the tower.
	339	+ name (str): The name of the tower.
	340	+ Returns:
	341	+ torch.Tensor: The output tensor of the tower network.
	342	+ """
	343	+ tower_units = list(map(int, self.params.tower_layers.strip().split(',')))
	344	+ y_tower = tower_input
	345	+ for tower_i, _ in enumerate(tower_units):
	346	+ tower_linear = nn.Linear(in_features=y_tower.shape[-1], out_features=tower_units[tower_i])
	347	+ tower_linear_out = tower_linear(y_tower)
	348	+ relu = nn.ReLU()
	349	+ y_tower = relu(tower_linear_out)
	350	+ return y_tower
	351	+
	352	+ def build_predictions(self, task_outputs: list) -> dict:
	353	+ """
	354	+ Build the predictions for the model.
	355	+
	356	+ Args:
	357	+ task_outputs (list): A list of task output tensors.
	358	+
	359	+ Returns:
	360	+ dict: A dictionary containing the predictions for ctr, cvr, and ctcvr.
	361	+ """
	362	+ preds = {}
	363	+ tower_outputs = {}
	364	+ for i, name in enumerate(self.tower_names):
	365	+ y = self.towers[name](task_outputs[i])
	366	+ y = self.towers_output_layers[name](y)
	367	+ y = torch.reshape(y, [-1, ])
	368	+ preds[name + '_logit'] = y
	369	+ preds[name] = torch.sigmoid(y)
	370	+ tower_outputs[name] = y
	371	+
	372	+ ctr_pred = preds.get('ctr')
	373	+ cvr_pred = preds.get('cvr')
	374	+ if ctr_pred is not None and cvr_pred is not None:
	375	+ if 'ctr_logit' not in preds or 'cvr_logit' not in preds:
	376	+ raise ValueError("Missing required keys in preds dictionary")
	377	+ else:
	378	+ preds['ctcvr'] = ctr_pred * cvr_pred
	379	+ preds['ctcvr_logit'] = preds['ctr_logit'] + preds['cvr_logit']
	380	+ return preds
	381	+
	382	+ def build_loss(self,
	383	+ labels: dict,
	384	+ y_ctr_logit: torch.Tensor,
	385	+ y_ctcvr_logit: torch.Tensor,
	386	+ ctr_weight,
	387	+ ctcvr_weight) -> torch.Tensor:
	388	+ """
	389	+ Build the loss function for the model.
	390	+
	391	+ Args:
	392	+ labels (dict): A dictionary containing the true labels for ctr and ctcvr.
	393	+ y_ctr_logit (torch.Tensor): The predicted ctr values.
	394	+ y_ctcvr_logit (torch.Tensor): The predicted ctcvr values.
	395	+ Returns:
	396	+ torch.Tensor: The combined loss tensor.
	397	+ """
	398	+ y_ctr_logit = y_ctr_logit.view(-1)
	399	+ y_ctcvr_logit = y_ctcvr_logit.view(-1)
	400	+ y = labels['y'].view(-1)
	401	+ z = labels['z'].view(-1)
	402	+
	403	+ bce_loss = nn.BCEWithLogitsLoss(pos_weight=torch.tensor([ctr_weight],
	404	+ dtype=torch.float32,
	405	+ device=y.device)
	406	+ )
	407	+ ctcvr_bce_loss = nn.BCEWithLogitsLoss(pos_weight=torch.tensor([ctcvr_weight],
	408	+ dtype=torch.float32,
	409	+ device=z.device)
	410	+ )
	411	+
	412	+ ctr_loss = bce_loss(y_ctr_logit, y)
	413	+ ctcvr_loss = ctcvr_bce_loss(y_ctcvr_logit, z)
	414	+ ctr_task_wgt = self.params.ctr_task_wgt
	415	+
	416	+ return ctr_task_wgt * ctr_loss + (1 - ctr_task_wgt) * ctcvr_loss
	417	+
	418	+ def build_optimizer(self):
	419	+ """
	420	+ Build the optimizer for training.
	421	+
	422	+ Args:
	423	+ loss (torch.Tensor): The loss tensor to minimize.
	424	+
	425	+ Returns:
	426	+ optim: The operation for applying gradients.
	427	+
	428	+ Raises:
	429	+ ValueError: If the optimizer type is not supported.
	430	+ """
	431	+ if self.params.optimizer == "Adam":
	432	+ optimizer = optim.Adam(
	433	+ params=self.parameters(),
	434	+ lr=self.params.learning_rate,
	435	+ betas=[0.9, 0.99], eps=1e-8
	436	+ )
	437	+ elif self.params.optimizer == "Adagrad":
	438	+ optimizer = optim.Adagrad(
	439	+ params=self.parameters(),
	440	+ lr=self.params.learning_rate,
	441	+ initial_accumulator_value=1e-6
	442	+ )
	443	+ elif self.params.optimizer == "Momentum":
	444	+ optimizer = optim.SGD(
	445	+ params=self.parameters(),
	446	+ lr=self.params.learning_rate,
	447	+ momentum=0.95
	448	+ )
	449	+ elif self.params.optimizer == "SGD":
	450	+ optimizer = optim.SGD(
	451	+ params=self.parameters(),
	452	+ lr=self.params.learning_rate,
	453	+ )
	454	+ else:
	455	+ raise ValueError("Unsupported optimizer type: {}".format(args.optimizer))
	456	+ return optimizer
	457	+
	458	+
	459	+def get_profile(profiler_start_step: int, profiler_end_step: int, profiling_save_path: str):
	460	+ warm_step = profiler_start_step
	461	+ active_step = profiler_end_step - warm_step +1
	462	+ print(f"[Profile INFO]: warm_step: {warm_step}, active_step: {active_step}, profiling_save_path: {profiling_save_path}")
	463	+ device = detect_device_type()
	464	+ if device == 'npu':
	465	+ import torch_npu
	466	+ g_prof_config = torch_npu.profiler._ExperimentalConfig(
	467	+ export_type=[
	468	+ torch_npu.profiler.ExportType.Text,
	469	+ torch_npu.profiler.ExportType.Db
	470	+ ],
	471	+ profiler_level=torch_npu.profiler.ProfilerLevel.Level2,
	472	+ msprof_tx=False,
	473	+ aic_metrics=torch_npu.profiler.AiCMetrics.AiCoreNone,
	474	+ l2_cache=False,
	475	+ op_attr=False,
	476	+ data_simplification=False,
	477	+ record_op_args=False,
	478	+ gc_detect_threshold=None)
	479	+
	480	+ return torch_npu.profiler.profile(
	481	+ activities=[
	482	+ torch_npu.profiler.ProfilerActivity.CPU,
	483	+ torch_npu.profiler.ProfilerActivity.NPU],
	484	+ schedule=torch_npu.profiler.schedule(wait=0, warmup=warm_step, active=active_step, repeat=1),
	485	+ on_trace_ready=torch_npu.profiler.tensorboard_trace_handler(profiling_save_path),
	486	+ record_shapes=False,
	487	+ profile_memory=False,
	488	+ with_stack=False,
	489	+ with_modules=False,
	490	+ with_flops=False,
	491	+ experimental_config=g_prof_config)
	492	+ elif device == 'cuda':
	493	+ from torch.profiler import profile, schedule, tensorboard_trace_handler, ProfilerActivity
	494	+ return profile(
	495	+ activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],
	496	+ schedule=schedule(wait=0, warmup=warm_step, active=active_step, repeat=1),
	497	+ on_trace_ready=tensorboard_trace_handler(profiling_save_path),
	498	+ record_shapes=False,
	499	+ profile_memory=False,
	500	+ )
	501	+ else:
	502	+ print(f"[Warning]: No supported acceleration device (CUDA/NPU) detected, profiler will be disabled")
	503	+ return None
	504	+
	505	+
	506	+def device_synchronize():
	507	+ if torch.cuda.is_available():
	508	+ torch.cuda.synchronize()
	509	+ elif torch.npu.is_available():
	510	+ torch.npu.synchronize()
	511	+
	512	+
	513	+def evaluate_op(model: TorchMmoeModel, te_files, device, args):
	514	+ te_dataset = TorchDataSet(te_files)
	515	+ te_ctr, te_ctcvr = pre_deal_dataset(te_dataset)
	516	+ test_dataloader = DataLoader(dataset=te_dataset,
	517	+ batch_size=args.batch_size,
	518	+ shuffle=True,
	519	+ collate_fn=collate_fn,
	520	+ prefetch_factor=100,
	521	+ num_workers=10)
	522	+ dl = iter(test_dataloader)
	523	+ input_sample, target_sample = next(dl)
	524	+ input_sample = {k: v.to(device) for k, v in input_sample.items()}
	525	+
	526	+ mod = 'compile' if args.enable_compile else 'eager'
	527	+ if mod == 'compile':
	528	+ model = torch.compile(model, mode='reduce-overhead', dynamic=False)
	529	+
	530	+ model.eval()
	531	+ with torch.no_grad():
	532	+ model(input_sample)
	533	+ device_synchronize()
	534	+
	535	+ prof = None
	536	+ if args.enable_profiler:
	537	+ profiling_save_path = args.profiler_save_path + '/' + MODEL_NAME + '/' + mod
	538	+ prof = get_profile(args.profiler_start_step, args.profiler_end_step, profiling_save_path)
	539	+ prof.start()
	540	+
	541	+ exec_times = []
	542	+ with torch.no_grad():
	543	+ for i in range(args.max_steps):
	544	+ e2etime1 = time.perf_counter()
	545	+ model(input_sample)
	546	+ device_synchronize()
	547	+ e2etime2 = time.perf_counter()
	548	+ elapsed_ms = (e2etime2 - e2etime1) * 1000
	549	+ print(f"iterations {i} / {args.max_steps}: [{mod}] e2e time: {elapsed_ms} ms.")
	550	+
	551	+ if i >= 10:
	552	+ exec_times.append(elapsed_ms)
	553	+
	554	+ if args.enable_profiler:
	555	+ prof.step()
	556	+ if args.enable_profiler:
	557	+ prof.stop()
	558	+ if exec_times:
	559	+ avg_time = sum(exec_times) / len(exec_times)
	560	+ print("Step time consumption statistics (excluding the first 10 steps)")
	561	+ print(f"\n>>> [{mod}] avg e2e: {avg_time:.3f} ms <<<")
	562	+
	563	+
	564	+def main(args):
	565	+ model = TorchMmoeModel(args)
	566	+ te_files = glob.glob("%stest/data_test.csv.hd5.*" % args.data_dir)
	567	+ device_type = detect_device_type()
	568	+ print(f"Using device type {device_type}")
	569	+ device = torch.device(device_type)
	570	+ model.to(device)
	571	+ evaluate_op(model, te_files, device, args)
	572	+
	573	+
	574	+if __name__ == "__main__":
	575	+ print(f"{MODEL_NAME} inference begin !!!")
	576	+ args = parse_arguments()
	577	+ main(args)

Abenchmarks/inference/README.md+96-0

@@ -0,0 +1,96 @@
	1	+# 推理说明
	2	+
	3	+本 README 说明如何使用 DCNv2，DIN，MMOE和ETA模型进行推理（含 eager / torch.compile 两种模式）。在运行完成后输出。
	4	+
	5	+---
	6	+
	7	+## 目录
	8	+
	9	+- [1. 数据集处理](#1-数据集处理)
	10	+ - [1.1 DIN and DCNv2](#11-din-and-dcnv2)
	11	+ - [1.2 MMOE and ETA](#12-mmoe-and-eta)
	12	+- [2. 模型推理](#2-模型推理)
	13	+ - [2.1 eager mode（默认）](#21-eager-mode默认)
	14	+ - [2.2 启用 torch.compile（可选）](#22-启用-torchcompile可选)
	15	+ - [2.3 采集 profile文件（可选）](#23-采集profile文件可选)
	16	+
	17	+---
	18	+
	19	+## 1. 数据集处理
	20	+
	21	+### 1.1 DIN and DCNv2
	22	+
	23	+DIN 和 DCNv2 的推理数据集来自 CTR_Algorithm 仓库示例数据，该数据集无需手动下载
	24	+
	25	+https://github.com/Prayforhanluo/CTR_Algorithm/tree/main/data
	26	+
	27	+在DIN或DCNv2目录下执行如下命令：
	28	+
	29	+```bash
	30	+git clone https://github.com/Prayforhanluo/CTR_Algorithm.git
	31	+```
	32	+
	33	+- 数据集目录：`CTR_Algorithm/data/data.csv`
	34	+
	35	+### 1.2 MMOE and ETA
	36	+
	37	+进入MMOE或ETA文件夹下，进行如下数据预处理
	38	+
	39	+推理数据集来自 https://tianchi.aliyun.com/dataset/408 Ali-CPP数据集，需做如下处理：
	40	+
	41	+1. 下载数据集，存放至当前目录下alicpp文件夹中
	42	+2. 后续流程参考：https://gitee.com/ascend/RecSDK/blob/develop_torch_benchmark/torch_examples_benchmark/model_zoo/README.md
	43	+
	44	+执行完成后数据集会默认生成至aliccp_out目录下。
	45	+
	46	+网络依赖已整理到 `requirements.txt`，可直接安装：
	47	+
	48	+```bash
	49	+pip install -r ./requirements.txt
	50	+```
	51	+
	52	+在模型执行时，默认加载数据集路径为: `./aliccp/aliccp_out/`；模型脚本同时支持运行时使用如下命令动态指定数据集目录：
	53	+
	54	+```bash
	55	+python eta.py --data_dir path/to/your/data/
	56	+```
	57	+
	58	+## 2. 模型推理
	59	+
	60	+进入DCNv2文件夹，直接执行推理脚本`dcnv2.py`即可
	61	+
	62	+- 注：运行该网络前需要先将patch文件加上，具体方式如下：
	63	+
	64	+ ```bash
	65	+ cd DCNv2
	66	+ unix2dos dcnv2.patch #如果在arm机器上需要执行
	67	+ git apply dcnv2.patch
	68	+ ```
	69	+
	70	+### 2.1 eager mode（默认）
	71	+
	72	+```bash
	73	+python dcnv2.py
	74	+```
	75	+
	76	+### 2.2 启用 `torch.compile`（可选）
	77	+
	78	+脚本已支持 `--enable_compile` 这样的开关，开启方式为：
	79	+
	80	+```bash
	81	+python dcnv2.py \
	82	+--enable_compile
	83	+```
	84	+
	85	+### 2.3 采集profile文件（可选）
	86	+
	87	+脚本已支持 `--enable_profiler` 这样的开关，开启方式为：
	88	+
	89	+```bash
	90	+python dcnv2.py \
	91	+ --enable_profiler \
	92	+ --profiler_start_step 5 \
	93	+ --profiler_end_step 6 \
	94	+```
	95	+
	96	+可以通过 `--profiler_start_step` 和 `--profiler_end_step` 分别设置profile开始和结束步数。

Abenchmarks/inference/requirements.txt+5-0

@@ -0,0 +1,5 @@
	1	+h5py==3.13.0
	2	+numpy==1.26.4
	3	+pandas
	4	+pytz
	5	+scikit-learn

Abenchmarks/llm/baichuan2-7B-Chat/README.md+124-0

@@ -0,0 +1,124 @@
	1	+# Baichuan2-7B-Chat 微调训练
	2	+
	3	+本 README 说明如何使用 Baichuan2-7B-Chat 模型权重，结合 LlamaFactory 提供的 `alpaca_zh_demo.json` 示例数据，完成数据下载、预处理与训练启动（含 eager / torch.compile 两种模式）。
	4	+
	5	+---
	6	+
	7	+## 目录
	8	+
	9	+- [1. 模型权重](#1-模型权重)
	10	+- [2. 数据获取](#2-数据获取)
	11	+- [3. 数据预处理](#3-数据预处理)
	12	+- [4. 模型训练](#4-模型训练)
	13	+ - [4.1 eager mode（默认）](#41-eager-mode默认)
	14	+ - [4.2 启用 torchcompile（可选）](#42-启用-torchcompile可选)
	15	+
	16	+---
	17	+
	18	+### 环境配置
	19	+
	20	+本项目依赖已整理到 `requirements.txt`，可直接安装：
	21	+
	22	+```bash
	23	+pip install -r ../utils/requirements.txt
	24	+```
	25	+
	26	+## 1. 模型权重
	27	+
	28	+- Hugging Face 模型：baichuan-inc/Baichuan2-7B-Chat
	29	+ https://huggingface.co/baichuan-inc/Baichuan2-7B-Chat
	30	+
	31	+- 可使用如下的自定义脚本下载
	32	+
	33	+```bash
	34	+python ../utils/download_hf.py --model baichuan-inc/Baichuan2-7B-Chat --save_path ./Baichuan2-7B-Chat
	35	+```
	36	+
	37	+## 2. 数据获取
	38	+
	39	+训练数据集来自 LlamaFactory 仓库示例数据：
	40	+
	41	+- `alpaca_zh_demo.json`
	42	+ https://github.com/hiyouga/LlamaFactory/blob/main/data/alpaca_zh_demo.json
	43	+
	44	+建议直接下载 raw 文件到本地：
	45	+
	46	+```bash
	47	+wget -O alpaca_zh_demo.json \
	48	+https://raw.githubusercontent.com/hiyouga/LlamaFactory/main/data/alpaca_zh_demo.json
	49	+```
	50	+
	51	+## 3. 数据预处理
	52	+
	53	+本项目提供 `preprocess.py`，用于把 Alpaca 格式（`instruction`/`input`/`output`）转换为 Baichuan 常用的 `conversations` 格式。
	54	+
	55	+### 命令
	56	+
	57	+```bash
	58	+python preprocess.py -i alpaca_zh_demo.json -o train.json
	59	+```
	60	+
	61	+## 4. 模型训练
	62	+
	63	+训练脚本：`run_baichuan2.sh`，脚本支持在 GPU 和 NPU 训练
	64	+
	65	+开始训练前，请修改脚本中的路径参数：
	66	+
	67	+- 模型权重路径（本地目录）
	68	+- 训练数据路径（预处理后的 `train.json`）
	69	+
	70	+### 4.1 eager mode（默认）
	71	+
	72	+```bash
	73	+bash run_baichuan2.sh
	74	+```
	75	+
	76	+### 4.2 启用 `torch.compile`（可选）
	77	+
	78	+可通过添加 `--enable_compile` 选项运行图模式
	79	+当在 GPU 上训练时，默认后端使用triton；当在 NPU 上训练时，可进一步指定后端为 mlir 或 dvm，默认使用 mlir。
	80	+
	81	+默认后端（mlir，可不写 --npu-backend）：
	82	+
	83	+```bash
	84	+bash run_baichuan2.sh \
	85	+ --enable_compile
	86	+```
	87	+
	88	+显式指定后端为 mlir：
	89	+
	90	+```bash
	91	+bash run_baichuan2.sh \
	92	+ --enable_compile \
	93	+ --npu-backend mlir
	94	+```
	95	+
	96	+切换后端为 dvm：
	97	+
	98	+```bash
	99	+bash run_baichuan2.sh \
	100	+ --enable_compile \
	101	+ --npu-backend dvm
	102	+```
	103	+
	104	+当在 NPU 上训练时，可通过 `--mfusion` 参数开启 MFusion 图算融合优化功能, 配合不同的NPU图模式后端, 进一步提升模型的性能，使用示例如下
	105	+
	106	+```bash
	107	+bash run_baichuan2.sh \
	108	+ --enable_compile \
	109	+ --npu-backend dvm \
	110	+ --mfusion
	111	+```
	112	+
	113	+### 4.3 采集profile文件（可选）
	114	+
	115	+脚本已支持 `--enable_profiler` 这样的开关，开启方式为：
	116	+
	117	+```bash
	118	+bash run_baichuan2.sh \
	119	+ --enable_profiler \
	120	+ --profiler_start_step 5 \
	121	+ --profiler_end_step 6 \
	122	+```
	123	+
	124	+可以通过 `--profiler_start_step` 和 `--profiler_end_step` 分别设置profile开始和结束步数。

Abenchmarks/llm/baichuan2-7B-Chat/preprocess.py+50-0

@@ -0,0 +1,50 @@
	1	+import json
	2	+import os
	3	+import argparse
	4	+
	5	+def convert_alpaca_to_baichuan(input_path, output_path):
	6	+ with open(input_path, "r", encoding="utf-8") as f:
	7	+ data = json.load(f)
	8	+
	9	+ converted_data = []
	10	+
	11	+ for item in data:
	12	+ instruction = item.get("instruction", "").strip()
	13	+ input_text = item.get("input", "").strip()
	14	+ output_text = item.get("output", "").strip()
	15	+
	16	+ if input_text:
	17	+ human_text = instruction + "\n" + input_text
	18	+ else:
	19	+ human_text = instruction
	20	+
	21	+ new_item = {
	22	+ "conversations": [
	23	+ {
	24	+ "from": "human",
	25	+ "value": human_text
	26	+ },
	27	+ {
	28	+ "from": "assistant",
	29	+ "value": output_text
	30	+ }
	31	+ ]
	32	+ }
	33	+
	34	+ converted_data.append(new_item)
	35	+
	36	+ with open(output_path, "w", encoding="utf-8") as f:
	37	+ json.dump(converted_data, f, ensure_ascii=False, indent=2)
	38	+
	39	+ print(f"Conversion completed, total {len(converted_data)} records")
	40	+ print(f"Output file: {output_path}")
	41	+
	42	+def build_argparser():
	43	+ p = argparse.ArgumentParser(description="Convert Alpaca JSON to Baichuan conversations JSON")
	44	+ p.add_argument("--input", "-i", required=True, help="Path to alpaca json file")
	45	+ p.add_argument("--output", "-o", required=True, help="Path to output json file")
	46	+ return p
	47	+
	48	+if __name__ == "__main__":
	49	+ args = build_argparser().parse_args()
	50	+ convert_alpaca_to_baichuan(args.input, args.output)

Abenchmarks/llm/baichuan2-7B-Chat/run_baichuan2.sh+23-0

@@ -0,0 +1,23 @@
	1	+#!/bin/bash
	2	+export TORCHINDUCTOR_CACHE_DIR="./cache"
	3	+export ASCEND_RT_VISIBLE_DEVICES=0
	4	+export CUDA_VISIBLE_DEVICES=0
	5	+export TORCH_COMPILE_DEBUG=1
	6	+export TORCH_NPU_USE_COMPATIBLE_IMPL=1
	7	+
	8	+rm -rf ./cache/*
	9	+mkdir -p ./cache logs
	10	+
	11	+python train_baichuan2_7B.py \
	12	+ --model_path $MODEL_PATH \
	13	+ --data_path $DATA_PATH \
	14	+ --output_dir "./baichuan2-finetuned" \
	15	+ --num_epochs 3 \
	16	+ --max_steps 200 \
	17	+ --batch_size 1 \
	18	+ --learning_rate 2e-5 \
	19	+ --max_length 1024 \
	20	+ --use_lora \
	21	+ --use_bf16 \
	22	+ --gradient_checkpointing \
	23	+ > logs/train_baichuan.log 2>&1

Abenchmarks/llm/baichuan2-7B-Chat/train_baichuan2_7B.py+308-0

@@ -0,0 +1,308 @@
	1	+import os
	2	+import json
	3	+import time
	4	+import sys
	5	+from pathlib import Path
	6	+import argparse
	7	+import logging
	8	+from typing import Dict, List, Optional
	9	+
	10	+import torch
	11	+from datasets import Dataset, load_dataset
	12	+
	13	+from transformers import (
	14	+ AutoTokenizer,
	15	+ AutoModelForCausalLM,
	16	+ TrainingArguments,
	17	+ Trainer,
	18	+ DataCollatorForSeq2Seq,
	19	+ TrainerCallback,
	20	+ set_seed,
	21	+)
	22	+from peft import LoraConfig, get_peft_model, TaskType
	23	+sys.path.append(str(Path(__file__).parent.parent))
	24	+from utils.utils import TimingCallback, get_profile, detect_device_type
	25	+
	26	+
	27	+logging.basicConfig(level=logging.INFO)
	28	+logger = logging.getLogger(__name__)
	29	+
	30	+MODEL_NAME = "Baichuan2-7B-Chat"
	31	+
	32	+class BaichuanTrainer(Trainer):
	33	+ def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=None):
	34	+ return super().compute_loss(model, inputs, return_outputs=return_outputs)
	35	+
	36	+
	37	+class Baichuan2Trainer:
	38	+ def __init__(self, args):
	39	+ self.args = args
	40	+ self.model = None
	41	+ self.tokenizer = None
	42	+ self.setup_training()
	43	+
	44	+ def setup_training(self):
	45	+ logger.info(f"Loading model/tokenizer from {self.args.model_path}")
	46	+
	47	+ self.tokenizer = AutoTokenizer.from_pretrained(
	48	+ self.args.model_path,
	49	+ trust_remote_code=True,
	50	+ padding_side="right",
	51	+ model_max_length=self.args.max_length,
	52	+ )
	53	+
	54	+ if self.tokenizer.pad_token is None:
	55	+ self.tokenizer.pad_token = self.tokenizer.eos_token
	56	+
	57	+ self.model = AutoModelForCausalLM.from_pretrained(
	58	+ self.args.model_path,
	59	+ trust_remote_code=True,
	60	+ torch_dtype=torch.bfloat16 if self.args.use_bf16 else torch.float16,
	61	+ use_cache=not self.args.gradient_checkpointing,
	62	+ )
	63	+
	64	+ if self.args.gradient_checkpointing:
	65	+ self.model.gradient_checkpointing_enable()
	66	+ self.model.config.use_cache = False
	67	+
	68	+ def apply_lora(self):
	69	+ if not self.args.use_lora:
	70	+ return
	71	+
	72	+ logger.info("Applying LoRA...")
	73	+
	74	+ target_modules = self.args.lora_target_modules.split(",") if self.args.lora_target_modules else ["W_pack"]
	75	+
	76	+ lora_config = LoraConfig(
	77	+ task_type=TaskType.CAUSAL_LM,
	78	+ r=self.args.lora_r,
	79	+ lora_alpha=self.args.lora_alpha,
	80	+ lora_dropout=self.args.lora_dropout,
	81	+ target_modules=target_modules,
	82	+ bias="none",
	83	+ )
	84	+ self.model = get_peft_model(self.model, lora_config)
	85	+ self.model.print_trainable_parameters()
	86	+
	87	+ def format_conversation(self, conversations: List[Dict]) -> str:
	88	+ text = ""
	89	+ for turn in conversations:
	90	+ role = turn.get("from", "")
	91	+ content = turn.get("value", "")
	92	+ if role == "human":
	93	+ text += f"<reserved_106>{content}"
	94	+ elif role == "assistant":
	95	+ text += f"<reserved_107>{content}"
	96	+ elif role == "system":
	97	+ text += f"<reserved_106>{content}"
	98	+ text += self.tokenizer.eos_token
	99	+ return text
	100	+
	101	+ def load_and_preprocess_data(self) -> Dataset:
	102	+ logger.info(f"Loading dataset from {self.args.data_path}")
	103	+
	104	+ if self.args.data_path.endswith(".json") or self.args.data_path.endswith(".jsonl"):
	105	+ with open(self.args.data_path, "r", encoding="utf-8") as f:
	106	+ if self.args.data_path.endswith(".jsonl"):
	107	+ raw = [json.loads(line) for line in f]
	108	+ else:
	109	+ raw = json.load(f)
	110	+
	111	+ formatted = []
	112	+ for item in raw:
	113	+ if "conversations" in item:
	114	+ text = self.format_conversation(item["conversations"])
	115	+ else:
	116	+ text = item.get("text", "")
	117	+ formatted.append({"text": text})
	118	+
	119	+ dataset = Dataset.from_list(formatted)
	120	+ else:
	121	+ dataset = load_dataset(self.args.data_path, split=self.args.split)
	122	+
	123	+ def preprocess_function(examples):
	124	+ tokenized = self.tokenizer(
	125	+ examples["text"],
	126	+ truncation=True,
	127	+ padding="max_length",
	128	+ max_length=self.args.max_length,
	129	+ return_attention_mask=True,
	130	+ )
	131	+ tokenized["labels"] = tokenized["input_ids"].copy()
	132	+ return tokenized
	133	+
	134	+ num_proc = self.args.num_proc if self.args.num_proc > 0 else None
	135	+ processed = dataset.map(
	136	+ preprocess_function,
	137	+ batched=True,
	138	+ remove_columns=dataset.column_names,
	139	+ num_proc=num_proc,
	140	+ load_from_cache_file=not self.args.overwrite_cache,
	141	+ )
	142	+ logger.info(f"Dataset size: {len(processed)}")
	143	+ return processed
	144	+
	145	+ def create_trainer(self, train_dataset, eval_dataset=None):
	146	+ has_eval = eval_dataset is not None
	147	+ eval_strategy = "steps" if has_eval else "no"
	148	+
	149	+ training_args = TrainingArguments(
	150	+ output_dir=self.args.output_dir,
	151	+ overwrite_output_dir=True,
	152	+ num_train_epochs=self.args.num_epochs,
	153	+ max_steps=self.args.max_steps,
	154	+ per_device_train_batch_size=self.args.batch_size,
	155	+ per_device_eval_batch_size=self.args.batch_size,
	156	+ gradient_accumulation_steps=self.args.gradient_accumulation_steps,
	157	+ learning_rate=self.args.learning_rate,
	158	+ weight_decay=self.args.weight_decay,
	159	+ warmup_ratio=self.args.warmup_ratio,
	160	+ lr_scheduler_type=self.args.lr_scheduler_type,
	161	+ logging_dir=os.path.join(self.args.output_dir, "logs"),
	162	+ logging_steps=self.args.logging_steps,
	163	+ save_steps=self.args.save_steps,
	164	+ save_total_limit=self.args.save_total_limit,
	165	+ eval_strategy=eval_strategy,
	166	+ eval_steps=self.args.eval_steps if has_eval else None,
	167	+ bf16=self.args.use_bf16,
	168	+ fp16=self.args.use_fp16,
	169	+ gradient_checkpointing=self.args.gradient_checkpointing,
	170	+ remove_unused_columns=False,
	171	+ report_to="none",
	172	+ dataloader_pin_memory=False,
	173	+ optim="adamw_torch",
	174	+ )
	175	+
	176	+ data_collator = DataCollatorForSeq2Seq(
	177	+ tokenizer=self.tokenizer,
	178	+ model=self.model,
	179	+ padding="max_length",
	180	+ max_length=self.args.max_length,
	181	+ label_pad_token_id=-100,
	182	+ )
	183	+
	184	+ mode = "compile" if self.args.enable_compile else "eager"
	185	+ prof=None
	186	+
	187	+ if self.args.enable_profiler:
	188	+ profiling_save_path = self.args.profiler_save_path + '/' + MODEL_NAME + '/' + mode
	189	+ prof = get_profile(self.args.profiler_start_step, self.args.profiler_end_step, profiling_save_path)
	190	+
	191	+ timing_callback = TimingCallback(prof, mode)
	192	+
	193	+ trainer = BaichuanTrainer(
	194	+ model=self.model,
	195	+ args=training_args,
	196	+ train_dataset=train_dataset,
	197	+ eval_dataset=eval_dataset,
	198	+ tokenizer=self.tokenizer,
	199	+ data_collator=data_collator,
	200	+ callbacks=[timing_callback],
	201	+ )
	202	+ return trainer
	203	+
	204	+ def train(self):
	205	+ logger.info("Starting training...")
	206	+
	207	+ if self.args.use_lora:
	208	+ self.apply_lora()
	209	+
	210	+ if self.args.enable_compile:
	211	+ logger.warning("Enable torch.compile (experimental on NPU).")
	212	+ try:
	213	+ self.model = torch.compile(self.model, dynamic=False)
	214	+ except Exception as e:
	215	+ logger.error(f"torch.compile failed: {e}. Continue without compile.")
	216	+
	217	+ dataset = self.load_and_preprocess_data()
	218	+
	219	+ trainer = self.create_trainer(dataset)
	220	+ train_result = trainer.train()
	221	+
	222	+ trainer.save_state()
	223	+ trainer.save_model(self.args.output_dir)
	224	+ self.tokenizer.save_pretrained(self.args.output_dir)
	225	+
	226	+ metrics = train_result.metrics
	227	+ trainer.log_metrics("train", metrics)
	228	+ trainer.save_metrics("train", metrics)
	229	+
	230	+ logger.info(f"Training completed! Model saved to {self.args.output_dir}")
	231	+ return metrics
	232	+
	233	+
	234	+def build_argparser():
	235	+ parser = argparse.ArgumentParser(description="Train Baichuan2-7B-Chat (LoRA) on NPU")
	236	+
	237	+ parser.add_argument("--npu-backend", type=str, default="mlir")
	238	+ parser.add_argument("--mfusion", action="store_true",
	239	+ help="Enable MFusion for graph fusion optimization")
	240	+ parser.add_argument("--model_path", type=str, required=True)
	241	+ parser.add_argument("--data_path", type=str, required=True)
	242	+ parser.add_argument("--output_dir", type=str, default="./baichuan2-finetuned")
	243	+ parser.add_argument("--num_epochs", type=int, default=3)
	244	+ parser.add_argument("--max_steps", type=int, default=-1)
	245	+ parser.add_argument("--batch_size", type=int, default=1)
	246	+ parser.add_argument("--gradient_accumulation_steps", type=int, default=1)
	247	+ parser.add_argument("--learning_rate", type=float, default=2e-5)
	248	+ parser.add_argument("--weight_decay", type=float, default=0.1)
	249	+ parser.add_argument("--warmup_ratio", type=float, default=0.03)
	250	+ parser.add_argument("--lr_scheduler_type", type=str, default="cosine")
	251	+ parser.add_argument("--max_length", type=int, default=1024)
	252	+ parser.add_argument("--gradient_checkpointing", action="store_true")
	253	+ parser.add_argument("--seed", type=int, default=42)
	254	+ parser.add_argument("--use_bf16", action="store_true")
	255	+ parser.add_argument("--use_fp16", action="store_true")
	256	+ parser.add_argument("--use_lora", action="store_true")
	257	+ parser.add_argument("--lora_r", type=int, default=8)
	258	+ parser.add_argument("--lora_alpha", type=int, default=32)
	259	+ parser.add_argument("--lora_dropout", type=float, default=0.1)
	260	+ parser.add_argument("--lora_target_modules", type=str, default="W_pack")
	261	+ parser.add_argument("--validation_split", type=float, default=0.0)
	262	+ parser.add_argument("--split", type=str, default="train")
	263	+ parser.add_argument("--num_proc", type=int, default=0)
	264	+ parser.add_argument("--overwrite_cache", action="store_true")
	265	+ parser.add_argument("--logging_steps", type=int, default=1)
	266	+ parser.add_argument("--save_steps", type=int, default=500)
	267	+ parser.add_argument("--save_total_limit", type=int, default=3)
	268	+ parser.add_argument("--eval_steps", type=int, default=500)
	269	+ parser.add_argument("--enable_compile", action="store_true")
	270	+ parser.add_argument("--enable_profiler", action="store_true")
	271	+ parser.add_argument("--profiler_save_path", type=str, default="./profile")
	272	+ parser.add_argument("--profiler_start_step", type=int, default=5,
	273	+ help="Output directory for trained model")
	274	+ parser.add_argument("--profiler_end_step", type=int, default=7,
	275	+ help="Output directory for trained model")
	276	+
	277	+ return parser
	278	+
	279	+
	280	+def main():
	281	+ args = build_argparser().parse_args()
	282	+ detect_device_type()
	283	+ os.environ['TORCHINDUCTOR_NPU_BACKEND']=args.npu_backend
	284	+ if args.npu_backend == "akg":
	285	+ os.environ['TORCHINDUCTOR_NPU_BACKEND'] = 'mlir'
	286	+ os.environ['TORCHINDUCTOR_USE_AKG'] = '1'
	287	+ if args.mfusion:
	288	+ os.environ['TORCHINDUCTOR_ENABLE_MFUSION']='1'
	289	+ os.makedirs(args.output_dir, exist_ok=True)
	290	+ os.makedirs(os.path.join(args.output_dir, "logs"), exist_ok=True)
	291	+
	292	+ set_seed(args.seed)
	293	+ torch.manual_seed(args.seed)
	294	+
	295	+ logger.info(f"npu train {MODEL_NAME}")
	296	+ trainer = Baichuan2Trainer(args)
	297	+ metrics = trainer.train()
	298	+
	299	+ if args.enable_compile:
	300	+ headers, values = torch._dynamo.utils.compile_times("csv")
	301	+ for header, value in zip(headers, values):
	302	+ if header == "PyCodeCache.load_by_key_path":
	303	+ numbers = [float(num.strip()) for num in value.split(',') if num.strip()]
	304	+ op_compile_time = sum(numbers)
	305	+ print(f"op_compile_time:{op_compile_time * 1e3} ms", )
	306	+
	307	+if __name__ == "__main__":
	308	+ main()

Abenchmarks/llm/glm4-9B-chat/README.md+106-0

@@ -0,0 +1,106 @@
	1	+# glm4-9B-chat 微调训练
	2	+
	3	+本 README 说明如何使用 glm4-9B-chat 模型权重，结合 LlamaFactory 提供的 `c4_demo.jsonl` 示例数据，完成数据下载、预处理与训练启动（含 eager / torch.compile 两种模式）。
	4	+
	5	+---
	6	+
	7	+## 目录
	8	+
	9	+- [1. 模型权重](#1-模型权重)
	10	+- [2. 数据获取](#2-数据获取)
	11	+- [3. 数据预处理](#3-数据预处理)
	12	+- [4. 模型训练](#4-模型训练)
	13	+ - [4.1 eager mode（默认）](#41-eager-mode默认)
	14	+ - [4.2 启用 torchcompile（可选）](#42-启用-torchcompile可选)
	15	+ - [4.3 采集profile文件（可选）](#43-采集profile文件可选)
	16	+
	17	+---
	18	+
	19	+## 1. 模型权重
	20	+
	21	+- Hugging Face 模型：glm4-9B-chat
	22	+ https://huggingface.co/zai-org/glm-4-9b-chat
	23	+
	24	+可在链接页面中 `Files and versions` 一栏直接下载。
	25	+
	26	+### 环境提示
	27	+
	28	+本项目依赖已整理到 `requirements.txt`，可直接安装：
	29	+
	30	+```bash
	31	+pip install -r ../utils/requirements.txt
	32	+```
	33	+
	34	+## 2. 数据获取
	35	+
	36	+训练数据集来自 LlamaFactory 仓库示例数据，可直接在如下连接获取：
	37	+
	38	+- `c4_demo.json`
	39	+ https://github.com/hiyouga/LlamaFactory/blob/main/data/c4_demo.jsonl
	40	+
	41	+或直接下载 raw 文件到本地：：
	42	+
	43	+```bash
	44	+wget -O c4_demo.json \
	45	+https://github.com/hiyouga/LlamaFactory/blob/main/data/c4_demo.jsonl
	46	+```
	47	+
	48	+## 3. 数据预处理
	49	+
	50	+在`train_glm4_9B.py`中已经对数据进行了处理
	51	+
	52	+## 4. 模型训练
	53	+
	54	+训练脚本：`train_glm4_9B.sh`，脚本支持在 GPU 和 NPU 训练
	55	+
	56	+开始训练前，请修改脚本中的路径参数：
	57	+
	58	+- 模型权重路径（本地目录）
	59	+- 训练数据路径（本地`c4_demo.json`）
	60	+
	61	+### 4.1 eager mode（默认）
	62	+
	63	+```bash
	64	+bash train_glm4_9B.sh
	65	+```
	66	+
	67	+### 4.2 启用 `torch.compile`（可选）
	68	+
	69	+可通过添加 `--enable_compile` 选项运行图模式
	70	+当在 GPU 上训练时，默认后端使用triton；当在 NPU 上训练时，可进一步指定后端为 mlir 或 dvm，默认使用 mlir。
	71	+
	72	+默认后端（mlir，可不写 --npu-backend）：
	73	+
	74	+```bash
	75	+bash train_glm4_9B.sh \
	76	+ --enable_compile
	77	+```
	78	+
	79	+显式指定后端为 mlir：
	80	+
	81	+```bash
	82	+bash train_glm4_9B.sh \
	83	+ --enable_compile \
	84	+ --npu-backend mlir
	85	+```
	86	+
	87	+切换后端为 dvm：：
	88	+
	89	+```bash
	90	+bash train_glm4_9B.sh \
	91	+ --enable_compile \
	92	+ --npu-backend dvm
	93	+```
	94	+
	95	+### 4.3 采集profile文件（可选）
	96	+
	97	+脚本已支持 `--enable_profiler` 这样的开关，开启方式为：
	98	+
	99	+```bash
	100	+bash train_glm4_9B.sh \
	101	+ --enable_profiler \
	102	+ --profiler_start_step 5 \
	103	+ --profiler_end_step 6 \
	104	+```
	105	+
	106	+可以通过 `--profiler_start_step` 和 `--profiler_end_step` 分别设置profile开始和结束步数。结束步数。

Abenchmarks/llm/glm4-9B-chat/run_glm4_9B.sh+20-0

@@ -0,0 +1,20 @@
	1	+export TORCHINDUCTOR_CACHE_DIR='./cache'
	2	+export ASCEND_RT_VISIBLE_DEVICES=0
	3	+export TORCH_COMPILE_DEBUG=1
	4	+export TORCH_NPU_USE_COMPATIBLE_IMPL=1
	5	+
	6	+rm -rf ./cache/*
	7	+mkdir -p ./cache logs
	8	+
	9	+python train_glm4_9B.py \
	10	+ --model_path "your model path" \
	11	+ --data_path "dataset path" \
	12	+ --output_dir "./glm4-finetuned" \
	13	+ --num_epochs 3 \
	14	+ --max_steps 200 \
	15	+ --batch_size 1 \
	16	+ --learning_rate 2e-5 \
	17	+ --max_length 512 \
	18	+ --use_lora \
	19	+ --use_bf16 \
	20	+ --pad_to_max_length > logs/train_glm4.log 2>&1

Abenchmarks/llm/glm4-9B-chat/train_glm4_9B.py+424-0

@@ -0,0 +1,424 @@
	1	+import torch
	2	+import os
	3	+import sys
	4	+from pathlib import Path
	5	+import torch.nn as nn
	6	+from datasets import Dataset, load_dataset
	7	+from transformers import (
	8	+ AutoModelForCausalLM,
	9	+ AutoTokenizer,
	10	+ TrainingArguments,
	11	+ Trainer,
	12	+ DataCollatorForLanguageModeling,
	13	+ BitsAndBytesConfig,
	14	+)
	15	+from peft import (
	16	+ LoraConfig,
	17	+ get_peft_model,
	18	+ prepare_model_for_kbit_training,
	19	+ TaskType
	20	+)
	21	+import json
	22	+from typing import Dict, Union, Any, Optional
	23	+import argparse
	24	+import logging
	25	+
	26	+
	27	+sys.path.append(str(Path(__file__).parent.parent))
	28	+from utils.utils import (
	29	+ TimingCallback,
	30	+ get_profile,
	31	+ detect_device_type
	32	+)
	33	+
	34	+logging.basicConfig(level=logging.INFO)
	35	+logger = logging.getLogger(__name__)
	36	+
	37	+model_name = 'GLM4-9B'
	38	+
	39	+class CustomTrainer(Trainer):
	40	+ def compute_loss(
	41	+ self,
	42	+ model: nn.Module,
	43	+ inputs: Dict[str, Union[torch.Tensor, Any]],
	44	+ return_outputs: bool = False,
	45	+ num_items_in_batch: Optional[torch.Tensor] = None,
	46	+ ):
	47	+ return super().compute_loss(
	48	+ model=model,
	49	+ inputs=inputs,
	50	+ return_outputs=return_outputs,
	51	+ num_items_in_batch=None
	52	+ )
	53	+
	54	+class GLM4Trainer:
	55	+ def __init__(self, args):
	56	+ self.args = args
	57	+ self.setup_training()
	58	+
	59	+ def setup_training(self):
	60	+ logger.info(f"Loading GLM-4 model from {self.args.model_path}")
	61	+
	62	+ bnb_config = None
	63	+ if self.args.use_4bit:
	64	+ bnb_config = BitsAndBytesConfig(
	65	+ load_in_4bit=True,
	66	+ bnb_4bit_quant_type="nf4",
	67	+ bnb_4bit_compute_dtype=torch.float16,
	68	+ bnb_4bit_use_double_quant=True
	69	+ )
	70	+
	71	+ self.model = AutoModelForCausalLM.from_pretrained(
	72	+ self.args.model_path,
	73	+ quantization_config=bnb_config if self.args.use_4bit else None,
	74	+ torch_dtype=torch.bfloat16 if self.args.use_bf16 else torch.float16,
	75	+ trust_remote_code=True,
	76	+ )
	77	+
	78	+ logger.info(f"Moving model to {self.args.device_type}...")
	79	+ self.model = self.model.to(self.args.device_type)
	80	+
	81	+ self.tokenizer = AutoTokenizer.from_pretrained(
	82	+ self.args.model_path,
	83	+ trust_remote_code=True
	84	+ )
	85	+
	86	+ if self.tokenizer.pad_token is None:
	87	+ self.tokenizer.pad_token = self.tokenizer.eos_token
	88	+ self.model.config.pad_token_id = self.tokenizer.eos_token_id
	89	+
	90	+ if self.args.gradient_checkpointing:
	91	+ self.model.gradient_checkpointing_enable()
	92	+ self.model.config.use_cache = False
	93	+
	94	+
	95	+ def apply_lora(self):
	96	+ if not self.args.use_lora:
	97	+ return
	98	+
	99	+ logger.info("Applying LoRA configuration...")
	100	+
	101	+ if self.args.use_4bit:
	102	+ self.model = prepare_model_for_kbit_training(self.model)
	103	+
	104	+ lora_config = LoraConfig(
	105	+ task_type=TaskType.CAUSAL_LM,
	106	+ r=self.args.lora_r,
	107	+ lora_alpha=self.args.lora_alpha,
	108	+ lora_dropout=self.args.lora_dropout,
	109	+ target_modules=self.get_lora_target_modules(),
	110	+ bias="none",
	111	+ )
	112	+
	113	+ self.model = get_peft_model(self.model, lora_config)
	114	+ self.model.print_trainable_parameters()
	115	+
	116	+ if self.args.enable_compile:
	117	+ self.model = torch.compile(self.model, dynamic=False)
	118	+
	119	+
	120	+ def get_lora_target_modules(self):
	121	+ target_modules = [
	122	+ "query_key_value",
	123	+ "dense",
	124	+ "dense_h_to_4h",
	125	+ "dense_4h_to_h",
	126	+ ]
	127	+
	128	+ model_modules = sorted([name for name, _ in self.model.named_modules()])
	129	+ available_modules = []
	130	+
	131	+ for target in target_modules:
	132	+ found = False
	133	+ for name in model_modules:
	134	+ if target in name:
	135	+ available_modules.append(name)
	136	+ found = True
	137	+ break
	138	+ if not found:
	139	+ logger.warning(f"Target module {target} not found in model")
	140	+
	141	+ available_modules = list(dict.fromkeys(available_modules))
	142	+
	143	+ if not available_modules:
	144	+ available_modules = ["qkv", "proj", "fc1", "fc2"]
	145	+ logger.warning(f"No avaliable, GLM-4 using LoRA target modules: {available_modules}")
	146	+ else:
	147	+ logger.info(f"GLM-4 LoRA target modules: {available_modules}")
	148	+
	149	+ return available_modules
	150	+
	151	+
	152	+ def load_and_preprocess_data(self) -> Dataset:
	153	+ logger.info(f"Loading dataset from {self.args.data_path}")
	154	+
	155	+ if self.args.data_path.endswith('.json') or self.args.data_path.endswith('.jsonl'):
	156	+ with open(self.args.data_path, 'r', encoding='utf-8') as f:
	157	+ if self.args.data_path.endswith('.jsonl'):
	158	+ data = [json.loads(line) for line in f]
	159	+ else:
	160	+ data = json.load(f)
	161	+
	162	+ formatted_data = []
	163	+ for item in data:
	164	+ if "conversations" in item or "messages" in item:
	165	+ messages = item.get("messages", item.get("conversations", []))
	166	+ text = self.tokenizer.apply_chat_template(
	167	+ messages,
	168	+ tokenize=False,
	169	+ add_generation_prompt=False
	170	+ )
	171	+ else:
	172	+ text = item.get("text", "")
	173	+
	174	+ formatted_data.append({"text": text})
	175	+
	176	+ dataset = Dataset.from_list(formatted_data)
	177	+ else:
	178	+ try:
	179	+ dataset = load_dataset(
	180	+ self.args.data_path,
	181	+ split=self.args.split
	182	+ )
	183	+ except:
	184	+ dataset = load_dataset(
	185	+ "json",
	186	+ data_files=self.args.data_path,
	187	+ split="train"
	188	+ )
	189	+
	190	+ tokenizer_ref = self.tokenizer
	191	+ max_length_ref = self.args.max_length
	192	+
	193	+ def preprocess_function(examples):
	194	+ tokenized = tokenizer_ref(
	195	+ examples["text"],
	196	+ truncation=True,
	197	+ padding="max_length" if self.args.pad_to_max_length else False,
	198	+ max_length=max_length_ref,
	199	+ return_tensors=None,
	200	+ return_attention_mask=True,
	201	+ )
	202	+ import copy
	203	+ tokenized["labels"] = copy.deepcopy(tokenized["input_ids"])
	204	+ return tokenized
	205	+
	206	+
	207	+ num_proc = self.args.num_proc if self.args.num_proc > 0 else None
	208	+ processed_dataset = dataset.map(
	209	+ preprocess_function,
	210	+ batched=True,
	211	+ remove_columns=dataset.column_names,
	212	+ num_proc=num_proc,
	213	+ load_from_cache_file=not self.args.overwrite_cache
	214	+ )
	215	+ logger.info(f"Dataset size: {len(processed_dataset)}")
	216	+ return processed_dataset
	217	+
	218	+
	219	+ def create_trainer(self, train_dataset, eval_dataset=None):
	220	+ has_eval = eval_dataset is not None
	221	+ eval_strategy = "steps" if has_eval else "no"
	222	+ save_strategy = "steps"
	223	+
	224	+ training_args = TrainingArguments(
	225	+ output_dir=self.args.output_dir,
	226	+ overwrite_output_dir=True,
	227	+ num_train_epochs=self.args.num_epochs,
	228	+ max_steps=self.args.max_steps if self.args.max_steps > 0 else 30,
	229	+ per_device_train_batch_size=self.args.batch_size,
	230	+ per_device_eval_batch_size=self.args.batch_size,
	231	+ gradient_accumulation_steps=self.args.gradient_accumulation_steps,
	232	+ weight_decay=self.args.weight_decay,
	233	+ logging_dir=f"{self.args.output_dir}/logs",
	234	+ logging_steps=self.args.logging_steps,
	235	+ save_steps=self.args.save_steps,
	236	+ save_total_limit=self.args.save_total_limit,
	237	+ eval_strategy=eval_strategy,
	238	+ eval_steps=self.args.eval_steps if has_eval else None,
	239	+ save_strategy=save_strategy,
	240	+ load_best_model_at_end=has_eval,
	241	+ metric_for_best_model="eval_loss",
	242	+ greater_is_better=False,
	243	+ learning_rate=self.args.learning_rate,
	244	+ lr_scheduler_type=self.args.lr_scheduler_type,
	245	+ fp16=self.args.use_fp16,
	246	+ bf16=self.args.use_bf16,
	247	+ gradient_checkpointing=self.args.gradient_checkpointing,
	248	+ report_to="none",
	249	+ ddp_find_unused_parameters=False if torch.cuda.device_count() > 1 else None,
	250	+ remove_unused_columns=False,
	251	+ dataloader_num_workers=self.args.dataloader_num_workers,
	252	+ group_by_length=self.args.group_by_length,
	253	+ length_column_name="length",
	254	+ prediction_loss_only=True,
	255	+ )
	256	+
	257	+ data_collator = DataCollatorForLanguageModeling(
	258	+ tokenizer=self.tokenizer,
	259	+ mlm=False,
	260	+ )
	261	+
	262	+ mod='compile' if self.args.enable_compile else 'eager'
	263	+ prof=None
	264	+
	265	+ if self.args.enable_profiler:
	266	+ profiling_save_path = self.args.profiler_save_path + '/' + model_name + '/' + mod
	267	+ prof = get_profile(self.args.profiler_start_step, self.args.profiler_end_step, profiling_save_path)
	268	+
	269	+ timing_callback = TimingCallback(prof, mod)
	270	+
	271	+ trainer = CustomTrainer(
	272	+ model=self.model,
	273	+ args=training_args,
	274	+ train_dataset=train_dataset,
	275	+ eval_dataset=eval_dataset,
	276	+ tokenizer=self.tokenizer,
	277	+ data_collator=data_collator,
	278	+ callbacks=[timing_callback],
	279	+ )
	280	+ return trainer
	281	+
	282	+
	283	+ def train(self):
	284	+ logger.info("Starting GLM-4 training...")
	285	+
	286	+ if self.args.use_lora:
	287	+ self.apply_lora()
	288	+
	289	+ dataset = self.load_and_preprocess_data()
	290	+
	291	+ if self.args.validation_split > 0:
	292	+ split_dataset = dataset.train_test_split(
	293	+ test_size=self.args.validation_split,
	294	+ seed=self.args.seed
	295	+ )
	296	+ train_dataset = split_dataset["train"]
	297	+ eval_dataset = split_dataset["test"]
	298	+ else:
	299	+ train_dataset = dataset
	300	+ eval_dataset = None
	301	+
	302	+ trainer = self.create_trainer(train_dataset, eval_dataset)
	303	+ train_result = trainer.train()
	304	+ if self.args.enable_compile:
	305	+ headers, values = torch._dynamo.utils.compile_times("csv")
	306	+ for header, value in zip(headers, values):
	307	+ if header == "PyCodeCache.load_by_key_path":
	308	+ numbers = [float(num.strip()) for num in value.split(',') if num.strip()]
	309	+ op_compile_time = sum(numbers)
	310	+ print(f"op_compile_time:{op_compile_time * 1e3} ms", )
	311	+
	312	+ trainer.save_model()
	313	+ self.tokenizer.save_pretrained(self.args.output_dir)
	314	+ metrics = train_result.metrics
	315	+ trainer.log_metrics("train", metrics)
	316	+ trainer.save_metrics("train", metrics)
	317	+ trainer.save_state()
	318	+ logger.info(f"Training completed! Model saved to {self.args.output_dir}")
	319	+ return metrics
	320	+
	321	+
	322	+def main():
	323	+ parser = argparse.ArgumentParser(description="Train GLM-4 model")
	324	+ parser.add_argument("--model_path", type=str, default="ZhipuAI/glm-4-9b-chat",
	325	+ help="Path to the pretrained GLM-4 model")
	326	+ parser.add_argument("--data_path", type=str, required=True,
	327	+ help="Path to training data (json/jsonl file or dataset name)")
	328	+ parser.add_argument("--output_dir", type=str, default="./glm4-finetuned",
	329	+ help="Output directory for trained model")
	330	+ parser.add_argument("--num_epochs", type=int, default=3,
	331	+ help="Number of training epochs")
	332	+ parser.add_argument("--batch_size", type=int, default=2,
	333	+ help="Batch size per device")
	334	+ parser.add_argument("--gradient_accumulation_steps", type=int, default=1,
	335	+ help="Gradient accumulation steps")
	336	+ parser.add_argument("--learning_rate", type=float, default=1e-4,
	337	+ help="Learning rate")
	338	+ parser.add_argument("--warmup_steps", type=int, default=50,
	339	+ help="Warmup steps")
	340	+ parser.add_argument("--max_steps", type=int, default=-1,
	341	+ help="Total training steps ")
	342	+ parser.add_argument("--weight_decay", type=float, default=0.01,
	343	+ help="Weight decay")
	344	+ parser.add_argument("--max_length", type=int, default=2048,
	345	+ help="Maximum sequence length")
	346	+ parser.add_argument("--lr_scheduler_type", type=str, default="cosine",
	347	+ choices=["linear", "cosine", "cosine_with_restarts", "constant"],
	348	+ help="Learning rate scheduler type")
	349	+ parser.add_argument("--use_lora", action="store_true",
	350	+ help="Use LoRA for parameter-efficient fine-tuning")
	351	+ parser.add_argument("--lora_r", type=int, default=8,
	352	+ help="LoRA rank")
	353	+ parser.add_argument("--lora_alpha", type=int, default=32,
	354	+ help="LoRA alpha")
	355	+ parser.add_argument("--lora_dropout", type=float, default=0.1,
	356	+ help="LoRA dropout")
	357	+ parser.add_argument("--use_4bit", action="store_true",
	358	+ help="Use 4-bit quantization")
	359	+ parser.add_argument("--use_fp16", action="store_true",
	360	+ help="Use FP16 precision")
	361	+ parser.add_argument("--use_bf16", action="store_true",
	362	+ help="Use BF16 precision")
	363	+ parser.add_argument("--gradient_checkpointing", action="store_true",
	364	+ help="Enable gradient checkpointing")
	365	+ parser.add_argument("--validation_split", type=float, default=0.1,
	366	+ help="Validation split ratio")
	367	+ parser.add_argument("--split", type=str, default="train",
	368	+ help="Dataset split to use")
	369	+ parser.add_argument("--num_proc", type=int, default=0,
	370	+ help="Number of processes for data preprocessing (0 = single process)")
	371	+ parser.add_argument("--pad_to_max_length", action="store_true",
	372	+ help="Pad sequences to max_length")
	373	+ parser.add_argument("--overwrite_cache", action="store_true",
	374	+ help="Overwrite cached features")
	375	+ parser.add_argument("--group_by_length", action="store_true",
	376	+ help="Group sequences by length for efficient training")
	377	+ parser.add_argument("--dataloader_num_workers", type=int, default=4,
	378	+ help="Number of workers for data loading")
	379	+ parser.add_argument("--seed", type=int, default=42,
	380	+ help="Random seed")
	381	+ parser.add_argument("--logging_steps", type=int, default=1,
	382	+ help="Log every X updates steps")
	383	+ parser.add_argument("--save_steps", type=int, default=500,
	384	+ help="Save checkpoint every X updates steps")
	385	+ parser.add_argument("--eval_steps", type=int, default=500,
	386	+ help="Evaluate every X updates steps")
	387	+ parser.add_argument("--save_total_limit", type=int, default=3,
	388	+ help="Limit the total amount of checkpoints")
	389	+ parser.add_argument("--report_to_tensorboard", action="store_true",
	390	+ help="Report metrics to TensorBoard")
	391	+ parser.add_argument("--enable_compile", action="store_true",
	392	+ help="Enable torch.compile and Inductor backend")
	393	+ parser.add_argument("--enable_profiler", action="store_true",
	394	+ help="Enable profiler for performance analysis")
	395	+ parser.add_argument("--profiler_start_step", type=int, default=5,
	396	+ help="Output directory for trained model")
	397	+ parser.add_argument("--profiler_end_step", type=int, default=8,
	398	+ help="Output directory for trained model")
	399	+ parser.add_argument("--profiler_save_path", type=str, default="./profile",
	400	+ help="Output directory for trained model")
	401	+ parser.add_argument("--npu-backend", type=str, default="mlir")
	402	+ parser.add_argument("--mfusion", action="store_true", help="Enable MFusion for graph fusion optimization")
	403	+ args = parser.parse_args()
	404	+ torch.manual_seed(args.seed)
	405	+ args.device_type = detect_device_type()
	406	+ os.environ['TORCHINDUCTOR_NPU_BACKEND']=args.npu_backend
	407	+ if args.npu_backend == "akg":
	408	+ os.environ['TORCHINDUCTOR_NPU_BACKEND'] = 'mlir'
	409	+ os.environ['TORCHINDUCTOR_USE_AKG'] = '1'
	410	+ if args.mfusion:
	411	+ os.environ['TORCHINDUCTOR_ENABLE_MFUSION']='1'
	412	+ print(f"{args.device_type} train {model_name}")
	413	+
	414	+ trainer = GLM4Trainer(args)
	415	+ metrics = trainer.train()
	416	+
	417	+ print("\n" + "="*50)
	418	+ print("GLM-4 Training completed successfully!")
	419	+ print(f"Model saved to: {args.output_dir}")
	420	+ print(f"Final training loss: {metrics.get('train_loss', 'N/A')}")
	421	+ print("="*50)
	422	+
	423	+if __name__ == "__main__":
	424	+ main()

Abenchmarks/llm/gpt-oss-20B/README.md+119-0

@@ -0,0 +1,119 @@
	1	+# gpt-oss-20B 微调训练
	2	+
	3	+本 README 说明如何使用 gpt-oss-20B 模型权重，结合 LlamaFactory 提供的 `c4_demo.jsonl` 示例数据，完成数据下载、预处理与训练启动（含 eager / torch.compile 两种模式）。
	4	+
	5	+---
	6	+
	7	+## 目录
	8	+
	9	+- [1. 模型权重](#1-模型权重)
	10	+- [2. 数据获取](#2-数据获取)
	11	+- [3. 数据预处理](#3-数据预处理)
	12	+- [4. 模型训练](#4-模型训练)
	13	+ - [4.1 eager mode（默认）](#41-eager-mode默认)
	14	+ - [4.2 启用 torchcompile（可选）](#42-启用-torchcompile可选)
	15	+ - [4.3 采集profile文件（可选）](#43-采集profile文件可选)
	16	+
	17	+---
	18	+
	19	+### 环境配置
	20	+
	21	+本项目依赖已整理到 `requirements.txt`，可直接安装：
	22	+
	23	+```bash
	24	+pip install -r ../utils/requirements.txt
	25	+```
	26	+
	27	+## 1. 模型权重
	28	+
	29	+- Hugging Face 模型地址：gpt-oss-20B
	30	+ https://huggingface.co/openai/gpt-oss-20b
	31	+
	32	+- 可使用如下的自定义脚本下载
	33	+
	34	+```bash
	35	+python ../utils/download_hf.py --model openai/gpt-oss-20b --save_path ./gpt-oss-20b
	36	+```
	37	+
	38	+## 2. 数据获取
	39	+
	40	+训练数据集来自 LlamaFactory 仓库示例数据，可直接在如下连接获取：
	41	+
	42	+- `c4_demo.json`
	43	+ https://github.com/hiyouga/LlamaFactory/blob/main/data/c4_demo.jsonl
	44	+
	45	+或直接下载 raw 文件到本地：
	46	+
	47	+```bash
	48	+wget -O c4_demo.json \
	49	+https://github.com/hiyouga/LlamaFactory/blob/main/data/c4_demo.jsonl
	50	+```
	51	+
	52	+## 3. 数据预处理
	53	+
	54	+在`train_gpt-oss_4B.py`中已经对数据进行了处理
	55	+
	56	+## 4. 模型训练
	57	+
	58	+训练脚本：`run_gpt-oss.sh`，脚本支持在 GPU 和 NPU 训练
	59	+
	60	+开始训练前，请修改脚本中的路径参数：
	61	+
	62	+- 模型权重路径
	63	+- 训练数据路径
	64	+
	65	+### 4.1 eager mode（默认）
	66	+
	67	+```bash
	68	+bash run_gpt-oss.sh
	69	+```
	70	+
	71	+### 4.2 启用 `torch.compile`（可选）
	72	+
	73	+可通过添加 `--enable_compile` 选项运行图模式
	74	+当在 GPU 上训练时，默认后端使用triton；当在 NPU 上训练时，可进一步指定后端为 mlir 或 dvm，默认使用 mlir。
	75	+
	76	+默认后端（mlir，可不写 --npu-backend）：
	77	+
	78	+```bash
	79	+bash run_gpt-oss.sh \
	80	+ --enable_compile
	81	+```
	82	+
	83	+显式指定后端为 mlir：
	84	+
	85	+```bash
	86	+bash run_gpt-oss.sh \
	87	+ --enable_compile \
	88	+ --npu-backend mlir
	89	+```
	90	+
	91	+切换后端为 dvm：
	92	+
	93	+```bash
	94	+bash run_gpt-oss.sh \
	95	+ --enable_compile \
	96	+ --npu-backend dvm
	97	+```
	98	+
	99	+当在 NPU 上训练时，可通过 `--mfusion` 参数开启 MFusion 图算融合优化功能, 配合不同的NPU图模式后端, 进一步提升模型的性能，使用示例如下
	100	+
	101	+```bash
	102	+bash run_gpt-oss.sh \
	103	+ --enable_compile \
	104	+ --npu-backend dvm \
	105	+ --mfusion
	106	+```
	107	+
	108	+### 4.3 采集profile文件（可选）
	109	+
	110	+脚本已支持 `--enable_profiler` 这样的开关，开启方式为：
	111	+
	112	+```bash
	113	+bash run_gpt-oss.sh \
	114	+ --enable_profiler \
	115	+ --profiler_start_step 5 \
	116	+ --profiler_end_step 6 \
	117	+```
	118	+
	119	+可以通过 `--profiler_start_step` 和 `--profiler_end_step` 分别设置profile开始和结束步数。

Abenchmarks/llm/gpt-oss-20B/run_gpt_oss.sh+23-0

@@ -0,0 +1,23 @@
	1	+#!/bin/bash
	2	+export TORCHINDUCTOR_CACHE_DIR="./cache"
	3	+export ASCEND_RT_VISIBLE_DEVICES=0
	4	+export CUDA_VISIBLE_DEVICES=0
	5	+export TORCH_COMPILE_DEBUG=1
	6	+export TORCH_NPU_USE_COMPATIBLE_IMPL=1
	7	+
	8	+rm -rf ./cache/*
	9	+mkdir -p ./cache logs
	10	+
	11	+python train_gpt_oss_20B.py \
	12	+ --model_path $MODEL_PATH \
	13	+ --data_path $DATA_PATH \
	14	+ --output_dir "./gpt-oss-finetuned" \
	15	+ --num_epochs 3 \
	16	+ --max_steps 200\
	17	+ --batch_size 1 \
	18	+ --learning_rate 2e-5 \
	19	+ --max_length 512 \
	20	+ --use_lora \
	21	+ --use_bf16 \
	22	+ --gradient_checkpointing \
	23	+ --pad_to_max_length > logs/train_gpt-oss.log 2>&1