add inference and llm model-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")

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

Abenchmarks/inference/README.md+92-0

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

@@ -0,0 +1,305 @@
	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.mfusion:
	285	+ os.environ['TORCHINDUCTOR_ENABLE_MFUSION']='1'
	286	+ os.makedirs(args.output_dir, exist_ok=True)
	287	+ os.makedirs(os.path.join(args.output_dir, "logs"), exist_ok=True)
	288	+
	289	+ set_seed(args.seed)
	290	+ torch.manual_seed(args.seed)
	291	+
	292	+ logger.info(f"npu train {MODEL_NAME}")
	293	+ trainer = Baichuan2Trainer(args)
	294	+ metrics = trainer.train()
	295	+
	296	+ if args.enable_compile:
	297	+ headers, values = torch._dynamo.utils.compile_times("csv")
	298	+ for header, value in zip(headers, values):
	299	+ if header == "PyCodeCache.load_by_key_path":
	300	+ numbers = [float(num.strip()) for num in value.split(',') if num.strip()]
	301	+ op_compile_time = sum(numbers)
	302	+ print(f"op_compile_time:{op_compile_time * 1e3} ms", )
	303	+
	304	+if __name__ == "__main__":
	305	+ 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

Abenchmarks/llm/gpt-oss-20B/train_gpt_oss_20B.py+390-0

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

Abenchmarks/llm/mamba-codestral-7B/README.md+118-0

@@ -0,0 +1,118 @@
	1	+# Mamba-Codestral-7B 微调训练
	2	+
	3	+本 README 说明如何使用 Mamba-Codestral-7B-v0.1 模型权重，结合 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	+
	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 模型：mistralai/Mamba-Codestral-7B-v0.1
	29	+ https://huggingface.co/mistralai/Mamba-Codestral-7B-v0.1
	30	+
	31	+- 可使用如下的自定义脚本下载
	32	+
	33	+```bash
	34	+python ../utils/download_hf.py --model mistralai/Mamba-Codestral-7B-v0.1 --save_path ./Mamba-Codestral-7B-v0.1
	35	+```
	36	+
	37	+## 2. 数据获取
	38	+
	39	+训练数据集来自 LlamaFactory 仓库示例数据，可直接在如下连接获取：
	40	+
	41	+- `c4_demo.json`
	42	+ https://github.com/hiyouga/LlamaFactory/blob/main/data/c4_demo.jsonl
	43	+
	44	+或直接下载 raw 文件到本地：
	45	+
	46	+```bash
	47	+wget -O c4_demo.json \
	48	+https://github.com/hiyouga/LlamaFactory/blob/main/data/c4_demo.jsonl
	49	+```
	50	+
	51	+## 3. 数据预处理
	52	+
	53	+在`train_mamba2_7B.py`中已经对数据进行了处理
	54	+
	55	+## 4. 模型训练
	56	+
	57	+训练脚本：`run_mamba2.sh`，脚本支持在 GPU 和 NPU 训练
	58	+
	59	+开始训练前，请修改脚本中的路径参数：
	60	+
	61	+- 模型权重路径
	62	+- 训练数据路径
	63	+
	64	+### 4.1 eager mode（默认）
	65	+
	66	+```bash
	67	+bash run_mamba2.sh
	68	+```
	69	+
	70	+### 4.2 启用 `torch.compile`（可选）
	71	+
	72	+可通过添加 `--enable_compile` 选项运行图模式
	73	+当在 GPU 上训练时，默认后端使用triton；当在 NPU 上训练时，可进一步指定后端为 mlir 或 dvm，默认使用 mlir。
	74	+
	75	+默认后端（mlir，可不写 --npu-backend）：
	76	+
	77	+```bash
	78	+bash run_mamba2.sh \
	79	+ --enable_compile
	80	+```
	81	+
	82	+显式指定后端为 mlir：
	83	+
	84	+```bash
	85	+bash run_mamba2.sh \
	86	+ --enable_compile \
	87	+ --npu-backend mlir
	88	+```
	89	+
	90	+切换后端为 dvm：
	91	+
	92	+```bash
	93	+bash run_mamba2.sh \
	94	+ --enable_compile \
	95	+ --npu-backend dvm
	96	+```
	97	+
	98	+当在 NPU 上训练时，可通过 `--mfusion` 参数开启 MFusion 图算融合优化功能, 配合不同的NPU图模式后端, 进一步提升模型的性能，使用示例如下
	99	+
	100	+```bash
	101	+bash run_mamba2.sh \
	102	+ --enable_compile \
	103	+ --npu-backend dvm \
	104	+ --mfusion
	105	+```
	106	+
	107	+### 4.3 采集profile文件（可选）
	108	+
	109	+脚本已支持 `--enable_profiler` 这样的开关，开启方式为：
	110	+
	111	+```bash
	112	+bash run_mamba2.sh \
	113	+ --enable_profiler \
	114	+ --profiler_start_step 5 \
	115	+ --profiler_end_step 6 \
	116	+```
	117	+
	118	+可以通过 `--profiler_start_step` 和 `--profiler_end_step` 分别设置profile开始和结束步数。

Abenchmarks/llm/mamba-codestral-7B/run_mamba2.sh+24-0

@@ -0,0 +1,24 @@
	1	+#!/bin/bash
	2	+export ASCEND_RT_VISIBLE_DEVICES=0
	3	+export CUDA_VISIBLE_DEVICES=0
	4	+export TORCHINDUCTOR_CACHE_DIR="./cache"
	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_mamba2_7B.py \
	12	+ --model_path $MODEL_PATH \
	13	+ --data_file $DATA_PATH \
	14	+ --output_dir "./mamba_codestral_lora" \
	15	+ --epochs 3 \
	16	+ --batch_size 1 \
	17	+ --gradient_accumulation_steps 1 \
	18	+ --learning_rate 1e-4 \
	19	+ --max_seq_length 256 \
	20	+ --lora_r 16 \
	21	+ --lora_alpha 32 \
	22	+ --lora_dropout 0.1 \
	23	+ --max_steps 200 \
	24	+ > logs/train_mamba2.log 2>&1

Abenchmarks/llm/mamba-codestral-7B/train_mamba2_7B.py+221-0

@@ -0,0 +1,221 @@
	1	+import os
	2	+import sys
	3	+import logging
	4	+from pathlib import Path
	5	+import argparse
	6	+import torch
	7	+import time
	8	+from datasets import Dataset as HFDataset
	9	+from transformers import (
	10	+ AutoTokenizer,
	11	+ AutoModelForCausalLM,
	12	+ TrainingArguments,
	13	+ Trainer,
	14	+ default_data_collator,
	15	+ TrainerCallback,
	16	+)
	17	+from peft import (
	18	+ LoraConfig,
	19	+ get_peft_model,
	20	+ TaskType,
	21	+)
	22	+sys.path.append(str(Path(__file__).parent.parent))
	23	+from utils.utils import TimingCallback, get_profile, detect_device_type
	24	+
	25	+logging.basicConfig(level=logging.INFO)
	26	+logger = logging.getLogger(__name__)
	27	+
	28	+MODEL_NAME = "Mamba2"
	29	+TARGET_MODULES = ["x_proj", "in_proj"]
	30	+
	31	+def patch_remove_ops_from_generate_list(op_names=None):
	32	+ try:
	33	+ import torch
	34	+ from torch_npu._inductor.ascend_npu_ir.ascend_npu_ir import config as anir_config
	35	+
	36	+ if not op_names:
	37	+ print("[patch] No op names provided, nothing to do.")
	38	+ return
	39	+
	40	+ for name in op_names:
	41	+ parts = name.split(".")
	42	+ op = torch.ops
	43	+ for p in parts:
	44	+ op = getattr(op, p)
	45	+
	46	+ if op in anir_config.GENERATE_LIST:
	47	+ anir_config.GENERATE_LIST.remove(op)
	48	+ print(f"[patch] Successfully removed {name} from GENERATE_LIST.")
	49	+ else:
	50	+ print(f"[patch] {name} not found in GENERATE_LIST (maybe already removed).")
	51	+
	52	+ except Exception as e:
	53	+ print(f"[patch] Failed to modify GENERATE_LIST: {e}")
	54	+
	55	+def setup_environment():
	56	+ os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"
	57	+
	58	+def parse_args():
	59	+ parser = argparse.ArgumentParser(description="Mamba Codestral LoRA finetune")
	60	+
	61	+ parser.add_argument("--model_path", type=str, default="/home/zhangyican/workspace/q4_data/Mamba-Codestral-7B-v0.1", help="Base model path")
	62	+ parser.add_argument("--data_file", type=str, default="./c4_demo.jsonl", help="Training data file path")
	63	+ parser.add_argument("--output_dir", type=str, default="./mamba_codestral_lora_no_quant", help="Output directory")
	64	+ parser.add_argument("--epochs", type=int, default=3, help="Number of training epochs")
	65	+ parser.add_argument("--batch_size", type=int, default=1, help="Batch size per device")
	66	+ parser.add_argument("--gradient_accumulation_steps", type=int, default=1, help="Gradient accumulation steps")
	67	+ parser.add_argument("--learning_rate", type=float, default=1e-4, help="Learning rate")
	68	+ parser.add_argument("--max_seq_length", type=int, default=256, help="Maximum sequence length")
	69	+ parser.add_argument("--lora_r", type=int, default=16, help="LoRA rank")
	70	+ parser.add_argument("--lora_alpha", type=int, default=32, help="LoRA alpha")
	71	+ parser.add_argument("--lora_dropout", type=float, default=0.1, help="LoRA dropout")
	72	+ parser.add_argument("--max_steps", type=int, default=100, help="Maximum training steps")
	73	+ parser.add_argument("--device", type=str, default="npu:0", help="NPU device, e.g., npu:0 / npu:6")
	74	+ parser.add_argument("--enable_compile", action="store_true", help="Whether to enable torch.compile")
	75	+ parser.add_argument("--enable_profiler", action="store_true")
	76	+ parser.add_argument("--profiler_save_path", type=str, default="./profile")
	77	+ parser.add_argument("--profiler_start_step", type=int, default=5,
	78	+ help="Start step for profiling")
	79	+ parser.add_argument("--profiler_end_step", type=int, default=6,
	80	+ help="End step for profiling")
	81	+ parser.add_argument("--npu-backend", type=str, default="mlir")
	82	+ parser.add_argument("--mfusion", action="store_true",
	83	+ help="Enable MFusion for graph fusion optimization")
	84	+
	85	+ return parser.parse_args()
	86	+
	87	+def main():
	88	+ args = parse_args()
	89	+ device = detect_device_type()
	90	+ os.environ['TORCHINDUCTOR_NPU_BACKEND']=args.npu_backend
	91	+ if args.mfusion:
	92	+ os.environ['TORCHINDUCTOR_ENABLE_MFUSION']='1'
	93	+ patch_remove_ops_from_generate_list(["aten.permute"])
	94	+ seed = 2
	95	+ torch.manual_seed(seed)
	96	+
	97	+ setup_environment()
	98	+ print("Environment setup complete")
	99	+
	100	+ print(f"Loading tokenizer from: {args.model_path}")
	101	+ tokenizer = AutoTokenizer.from_pretrained(
	102	+ args.model_path,
	103	+ trust_remote_code=True,
	104	+ use_fast=True
	105	+ )
	106	+ if tokenizer.pad_token is None:
	107	+ tokenizer.pad_token = tokenizer.eos_token
	108	+
	109	+ print(f"Loading base model from: {args.model_path} (no quantization)")
	110	+ base_model = AutoModelForCausalLM.from_pretrained(
	111	+ args.model_path,
	112	+ torch_dtype=torch.bfloat16, # Use bfloat16 precision
	113	+ trust_remote_code=True,
	114	+ use_cache=False, # Disable cache to save memory
	115	+ )
	116	+
	117	+ base_model.gradient_checkpointing_enable()
	118	+
	119	+ print(f"Configuring LoRA, target modules: {TARGET_MODULES}")
	120	+ peft_config = LoraConfig(
	121	+ task_type=TaskType.CAUSAL_LM,
	122	+ r=args.lora_r,
	123	+ lora_alpha=args.lora_alpha,
	124	+ lora_dropout=args.lora_dropout,
	125	+ target_modules=TARGET_MODULES,
	126	+ bias="none",
	127	+ )
	128	+
	129	+ print("Attaching LoRA adapters to model...")
	130	+ peft_model = get_peft_model(base_model, peft_config)
	131	+ peft_model.print_trainable_parameters() # Print number of trainable parameters
	132	+
	133	+ peft_model.to(device)
	134	+
	135	+ print("Preparing dataset...")
	136	+ from datasets import load_dataset
	137	+ dataset = load_dataset(
	138	+ "json",
	139	+ data_files=args.data_file,
	140	+ split="train"
	141	+ )
	142	+
	143	+ def tokenize_function(examples):
	144	+ return tokenizer(
	145	+ examples["text"],
	146	+ truncation=True,
	147	+ padding="max_length",
	148	+ max_length=args.max_seq_length,
	149	+ return_tensors=None
	150	+ )
	151	+
	152	+ train_dataset = dataset.map(tokenize_function, batched=True, remove_columns=["text"])
	153	+ train_dataset = train_dataset.map(lambda x: {"labels": x["input_ids"]}, batched=True)
	154	+
	155	+ if args.enable_compile:
	156	+ peft_model = torch.compile(
	157	+ peft_model,
	158	+ dynamic=False
	159	+ )
	160	+
	161	+ training_args = TrainingArguments(
	162	+ output_dir=args.output_dir,
	163	+ num_train_epochs=args.epochs,
	164	+ per_device_train_batch_size=args.batch_size,
	165	+ gradient_accumulation_steps=args.gradient_accumulation_steps,
	166	+ gradient_checkpointing=True,
	167	+ gradient_checkpointing_kwargs={"use_reentrant": False},
	168	+ logging_steps=1,
	169	+ max_steps=args.max_steps,
	170	+ learning_rate=args.learning_rate,
	171	+ warmup_ratio=0.1,
	172	+ lr_scheduler_type="cosine",
	173	+ bf16=True,
	174	+ fp16=False,
	175	+ save_strategy="steps",
	176	+ save_total_limit=2,
	177	+ remove_unused_columns=False,
	178	+ report_to="none",
	179	+ dataloader_pin_memory=False,
	180	+ max_grad_norm=0.3,
	181	+ logging_dir=f"{args.output_dir}/logs",
	182	+ )
	183	+
	184	+ mode = "compile" if args.enable_compile else "eager"
	185	+ prof=None
	186	+
	187	+ if args.enable_profiler:
	188	+ profiling_save_path = args.profiler_save_path + '/' + MODEL_NAME + '/' + mode
	189	+ prof = get_profile(args.profiler_start_step, args.profiler_end_step, profiling_save_path)
	190	+
	191	+ timing_callback = TimingCallback(prof, mode)
	192	+
	193	+ print("Initializing Trainer...")
	194	+ trainer = Trainer(
	195	+ model=peft_model,
	196	+ args=training_args,
	197	+ train_dataset=train_dataset,
	198	+ tokenizer=tokenizer,
	199	+ data_collator=default_data_collator,
	200	+ callbacks=[timing_callback],
	201	+ )
	202	+
	203	+ print("Starting LoRA training...")
	204	+ trainer.train()
	205	+
	206	+ if args.enable_compile:
	207	+ headers, values = torch._dynamo.utils.compile_times("csv")
	208	+ for header, value in zip(headers, values):
	209	+ if header == "PyCodeCache.load_by_key_path":
	210	+ numbers = [float(num.strip()) for num in value.split(',') if num.strip()]
	211	+ op_compile_time = sum(numbers)
	212	+ print(f"op_compile_time:{op_compile_time * 1e3} ms", )
	213	+
	214	+ print(f"Saving LoRA weights to {args.output_dir}")
	215	+ peft_model.save_pretrained(args.output_dir)
	216	+ tokenizer.save_pretrained(args.output_dir)
	217	+
	218	+ print("Training complete!")
	219	+
	220	+if __name__ == "__main__":
	221	+ main()

Abenchmarks/llm/sdxl-base-1.0/README.md+114-0

Abenchmarks/llm/sdxl-base-1.0/run_sdxl.sh+21-0

Abenchmarks/llm/sdxl-base-1.0/train_sdxl_base.py+274-0

Abenchmarks/llm/utils/download_hf.py+52-0

Abenchmarks/llm/utils/requirements.txt+15-0

Abenchmarks/llm/utils/utils.py+135-0