QLoRA特性支持-MindSpeed-LLM-AtomGit

Mconvert_ckpt.py+2-0

	@@ -82,6 +82,8 @@ def main():
82	help='model type of huggingface')	82	help='model type of huggingface')
83	parser.add_argument('--ckpt-cfg-path', type=str, default="configs/checkpoint/model_cfg.json",	83	parser.add_argument('--ckpt-cfg-path', type=str, default="configs/checkpoint/model_cfg.json",
84	help="Path to the config directory. If not specified, the default path in the repository will be used.")	84	help="Path to the config directory. If not specified, the default path in the repository will be used.")
		85	+ parser.add_argument('--qlora-nf4', action='store_true',
		86	+ help='use bitsandbytes nf4 to quantize model.')
85	parser.add_argument('--orm', action="store_true", default=False,	87	parser.add_argument('--orm', action="store_true", default=False,
86	help='Specify the ORM ckpt conversion, convert additional rm_head layer in ORM.')	88	help='Specify the ORM ckpt conversion, convert additional rm_head layer in ORM.')
87	known_args, _ = parser.parse_known_args()	89	known_args, _ = parser.parse_known_args()

Adocs/features/qlora.md+90-0

@@ -0,0 +1,90 @@
	1	+# QLoRA
	2	+
	3	+![QLoRA](../../sources/images/qlora/qlora.png)
	4	+
	5	+## 特性介绍
	6	+
	7	+### LoRA
	8	+
	9	+模型中的线性层通常是将激活值$x$与权重矩阵的转置$W^T$进行矩阵乘运算，并加上$bias$：
	10	+
	11	+$$
	12	+y = xW^T + bias
	13	+$$
	14	+
	15	+矩阵的秩（Rank）揭示了这个矩阵的“信息量”，预训练模型中的这些线性层的权重矩阵通常是满秩的。
	16	+
	17	+对于预训练的权重矩阵$W_0\in \mathbb{R}^{n\times m}$，可以认为它在微调阶段更新后变成了$W_0 + \Delta W = W_0 + AB$，其中$A\in \mathbb{R}^{n\times r}, B\in \mathbb{R}^{r\times m}$，那么矩阵$\Delta W = AB\in \mathbb{R}^{n\times m}$，且$rank(\Delta W)\in[0, min(n, m, r)]$，我们选取$r\ll min(n, m)$，所以$rank(\Delta W)\in[0, r]$。
	18	+
	19	+这样冻结$W_0$部分权重，仅更新参数量极少的旁路矩阵$A,B$的训练是非常高效的，很大程度上减少了低秩矩阵中的冗余信息占用的空间和计算量。
	20	+
	21	+在前向传播时，增加了LoRA的线性层（矩阵部分）变成了计算：
	22	+
	23	+$$
	24	+y = xW_0^T + x A^T B^T
	25	+$$
	26	+
	27	+反向传播的输入是Loss对当前层的输出 $y$ 的梯度 $\frac{\partial \mathcal{L}}{\partial y}$，需要更新旁路矩阵 $A, B$ 的参数，所以需要计算 $A, B$ 的梯度 $\frac{\partial \mathcal{L}}{\partial A}, \frac{\partial \mathcal{L}}{\partial B}$，仅对$A$或$B$求偏导可以将 $xW_0^T$ 看做常数项，这个过程不涉及$W_0$，但为了继续反向传播，当前层还需要返回 $\mathcal{L}$ 对输入的的梯度：
	28	+
	29	+$$
	30	+\frac{\partial \mathcal{L}}{\partial x}=\frac{\partial \mathcal{L}}{\partial y} \frac{\partial y}{\partial x} = \frac{\partial \mathcal{L}}{\partial y} W_0 + \frac{\partial \mathcal{L}}{\partial y} BA
	31	+$$
	32	+
	33	+因此，$W_0$的参数不需要更新，但是仍需要在前向和反向时各参与一次矩阵乘运算。
	34	+
	35	+### QLoRA
	36	+
	37	+#### 线性层量化
	38	+
	39	+量化（Quantization）技术可以压缩模型大小的同时保持较高的精度，无校准集的仅权重量化可以使
	40	+
	41	+$$
	42	+\phi \big(W - Q^{-1}(Q(W))\big)
	43	+$$
	44	+
	45	+尽可能小，误差在可以接受的范围内，其中$W$表示浮点权重，$Q()$和$Q^{-1}()$分别代表量化和反量化函数，$\phi$代表某种损失函数比如平均绝对误差等。
	46	+
	47	+我们使用NF4量化来进行权重量化，这也是开源社区流行的QLoRA量化方式，该量化算法被实现在[bitsandbytes](https://github.com/bitsandbytes-foundation/bitsandbytes)，它在开源社区广受欢迎。
	48	+
	49	+> 我们已经将支持NPU硬件的NF4量化功能贡献到bitsandbytes多硬件后端重构分支，但由于bitsandbytes官方还未正式将该分支在PyPi上发布，本仓库暂使用NPU版本的bitsandbytes，可以通过`pip install bitsandbytes-npu-beta`来安装。
	50	+
	51	+NF4是一种查找表量化，NF4表接近高斯分布下的数据的最佳表示方式，我们可以在QLoRA微调前先对权重$W_0$进行量化：
	52	+
	53	+$$
	54	+W_0^{NF4} = Q(W_0)
	55	+$$
	56	+
	57	+然后在微调阶段每次需要$W_0$参与运算时对其进行反量化，即在前向传播中计算：
	58	+
	59	+$$
	60	+y = x\big(Q^{-1}(W_0^{NF4})\big)^T + x A^T B^T
	61	+$$
	62	+
	63	+在反向传播中计算：
	64	+
	65	+$$
	66	+\frac{\partial \mathcal{L}}{\partial x} = \frac{\partial \mathcal{L}}{\partial y} Q^{-1}(W_0^{NF4}) + \frac{\partial \mathcal{L}}{\partial y} BA
	67	+$$
	68	+
	69	+QLoRA在LoRA的基础上，对主干部分的权重进行量化，大幅降低了LoRA微调时的显存占用。
	70	+
	71	+## 使用方法
	72	+
	73	+1、将hf权重转换为mcore权重时，增加`--qlora-nf4`选项开启QLoRA的NF4量化，目前不支持其它量化方式；
	74	+
	75	+2、在微调时，增加`--qlora`使用QLoRA微调，开启时请确保配置的权重路径是NF4量化后的mcore权重。
	76	+
	77	+## 使用效果
	78	+
	79	+由于仅对线性层的权重矩阵进行量化，激活值、梯度、优化器状态等仍然保持原精度，因此使用QLoRA可以给整体带来的内存收益比例会受到batch size、数据序列长度等因素的影响。
	80	+
	81	+作为参考，我们在Llama-2-70b和Mixtral-8x7b上进行量化测试：
	82	+
	83	+\| 模型 \| 原始权重大小 \| NF4量化后权重大小 \| 节省内存 \|
	84	+\| ------------ \| ------------ \| ----------------- \| -------- \|
	85	+\| Llama-2-70b \| 129GB \| 35GB \| 94GB \|
	86	+\| Mixtral-8x7b \| 87GB \| 24GB \| 63GB \|
	87	+
	88	+均可使用64GB单卡进行QLoRA微调，实测Llama-2-70b单batch微调过程中NPU内存占用约50GB左右。
	89	+
	90	+> QLoRA支持分布式LoRA、lora-fusion、PP、TP等LoRA支持的特性，并且精度正常，更多特性的亲和性还在补充验证中。

Aexamples/mcore/llama2/ckpt_convert_llama2_hf2mcore_qlora.sh+18-0

@@ -0,0 +1,18 @@
	1	+export CUDA_DEVICE_MAX_CONNECTIONS=1
	2	+
	3	+# 修改 ascend-toolkit 路径
	4	+source /usr/local/Ascend/ascend-toolkit/set_env.sh
	5	+
	6	+# 开启 --qlora-nf4 选项使用 QLoRA
	7	+python convert_ckpt.py \
	8	+ --model-type GPT \
	9	+ --load-model-type hf \
	10	+ --save-model-type mg \
	11	+ --target-tensor-parallel-size 1 \
	12	+ --target-pipeline-parallel-size 1\
	13	+ --load-dir ./model_from_hf/llama-2-70b-hf/ \
	14	+ --save-dir ./model_weights/Llama2-mcore/ \
	15	+ --tokenizer-model ./model_from_hf/llama-2-70b-hf/tokenizer.model \
	16	+ --use-mcore-models \
	17	+ --qlora-nf4 \
	18	+ --model-type-hf llama2

Aexamples/mcore/llama2/tune_llama2_70b_qlora_ptd.sh+116-0

@@ -0,0 +1,116 @@
	1	+#!/bin/bash
	2	+
	3	+export CUDA_DEVICE_MAX_CONNECTIONS=1
	4	+export PYTORCH_NPU_ALLOC_CONF=expandable_segments:True
	5	+export HCCL_CONNECT_TIMEOUT=1200
	6	+export HCCL_EXEC_TIMEOUT=1200
	7	+
	8	+NPUS_PER_NODE=1
	9	+MASTER_ADDR=localhost
	10	+MASTER_PORT=6080
	11	+NNODES=1
	12	+NODE_RANK=0
	13	+WORLD_SIZE=$(($NPUS_PER_NODE*$NNODES))
	14	+
	15	+CKPT_SAVE_DIR="your model save ckpt path"
	16	+DATA_PATH="your data path"
	17	+TOKENIZER_MODEL="your tokenizer path"
	18	+CKPT_LOAD_DIR="your model ckpt path"
	19	+
	20	+TP=1
	21	+PP=1
	22	+
	23	+DISTRIBUTED_ARGS="
	24	+ --nproc_per_node $NPUS_PER_NODE \
	25	+ --nnodes $NNODES \
	26	+ --node_rank $NODE_RANK \
	27	+ --master_addr $MASTER_ADDR \
	28	+ --master_port $MASTER_PORT \
	29	+"
	30	+
	31	+
	32	+TRAINING_ARGS="
	33	+ --tensor-model-parallel-size ${TP} \
	34	+ --pipeline-model-parallel-size ${PP} \
	35	+ --sequence-parallel \
	36	+ --use-mcore-models \
	37	+ --prompt-type llama2 \
	38	+ --tokenizer-type PretrainedFromHF \
	39	+ --tokenizer-name-or-path ${TOKENIZER_MODEL} \
	40	+ --micro-batch-size 1 \
	41	+ --global-batch-size 1 \
	42	+ --train-iters 2000 \
	43	+ --lr 1.0e-6 \
	44	+ --lr-decay-style cosine \
	45	+ --lr-warmup-fraction 0.00 \
	46	+ --weight-decay 0.0 \
	47	+ --clip-grad 1.0 \
	48	+ --swiglu \
	49	+ --position-embedding-type rope \
	50	+ --untie-embeddings-and-output-weights \
	51	+ --disable-bias-linear \
	52	+ --hidden-dropout 0.0 \
	53	+ --normalization RMSNorm \
	54	+ --use-flash-attn \
	55	+ --no-masked-softmax-fusion \
	56	+ --attention-softmax-in-fp32 \
	57	+ --init-method-std 0.01 \
	58	+ --initial-loss-scale 1 \
	59	+ --adam-beta1 0.9 \
	60	+ --adam-beta2 0.999 \
	61	+ --adam-eps 1e-8 \
	62	+ --no-gradient-accumulation-fusion \
	63	+ --no-load-optim \
	64	+ --no-load-rng \
	65	+ --variable-seq-lengths \
	66	+ --rotary-base 10000 \
	67	+ --norm-epsilon 1e-05 \
	68	+ --make-vocab-size-divisible-by 1 \
	69	+ --padded-vocab-size 32000 \
	70	+ --vocab-size 32000 \
	71	+ --bf16 \
	72	+ --num-layers 80 \
	73	+ --hidden-size 8192 \
	74	+ --ffn-hidden-size 28672 \
	75	+ --num-attention-heads 64 \
	76	+ --seq-length 4096 \
	77	+ --max-position-embeddings 4096 \
	78	+ --use-rotary-position-embeddings \
	79	+ --group-query-attention \
	80	+ --num-query-groups 8 \
	81	+ --tokenizer-not-use-fast \
	82	+ --attention-dropout 0.0 \
	83	+"
	84	+
	85	+FINETUNE_ARGS="
	86	+ --finetune \
	87	+ --stage sft \
	88	+ --is-instruction-dataset \
	89	+ --lora-r 16 \
	90	+ --lora-alpha 32 \
	91	+ --lora-target-modules linear_qkv linear_proj linear_fc1 linear_fc2 \
	92	+ --qlora \
	93	+"
	94	+
	95	+DATA_ARGS="
	96	+ --data-path $DATA_PATH \
	97	+ --split 100,0,0 \
	98	+"
	99	+
	100	+OUTPUT_ARGS="
	101	+ --log-interval 1 \
	102	+ --save-interval 2000 \
	103	+ --eval-interval 10000 \
	104	+ --eval-iters 0 \
	105	+"
	106	+
	107	+torchrun $DISTRIBUTED_ARGS posttrain_gpt.py \
	108	+ $TRAINING_ARGS \
	109	+ $FINETUNE_ARGS \
	110	+ $DATA_ARGS \
	111	+ $OUTPUT_ARGS \
	112	+ --load $CKPT_LOAD_DIR \
	113	+ --save $CKPT_SAVE_DIR \
	114	+ --log-throughput \
	115	+ --distributed-backend nccl \
	116	+ \| tee logs/tune_llama2_70b_mocre_qlora.log

Aexamples/mcore/mixtral/ckpt_convert_mixtral_hf2mcore_qlora.sh+20-0

@@ -0,0 +1,20 @@
	1	+export CUDA_DEVICE_MAX_CONNECTIONS=1
	2	+
	3	+# 修改 ascend-toolkit 路径
	4	+source /usr/local/Ascend/ascend-toolkit/set_env.sh
	5	+
	6	+# 设置需要的并行配置
	7	+python convert_ckpt.py \
	8	+ --model-type GPT \
	9	+ --load-model-type hf \
	10	+ --save-model-type mg \
	11	+ --params-dtype bf16 \
	12	+ --qlora-nf4 \
	13	+ --target-tensor-parallel-size 1 \
	14	+ --target-pipeline-parallel-size 1 \
	15	+ --target-expert-parallel-size 1 \
	16	+ --load-dir ./model_from_hf/Mixtral-hf/ \
	17	+ --save-dir ./model_weights/Mixtral-mcore/ \
	18	+ --tokenizer-model ./model_from_hf/Mixtral-hf/tokenizer.model \
	19	+ --use-mcore-models \
	20	+ --model-type-hf mixtral

Mmindspeed_llm/tasks/checkpoint/saver.py+23-0

	@@ -20,6 +20,7 @@ import sys
20		20
21	import torch	21	import torch
22	from megatron.core import mpu	22	from megatron.core import mpu
		23	+import megatron.core.tensor_parallel.layers as tpl
23	from megatron.training.checkpointing import save_checkpoint	24	from megatron.training.checkpointing import save_checkpoint
24		25
25	from .models import get_megatron_model	26	from .models import get_megatron_model
	@@ -27,6 +28,11 @@ from .models import get_megatron_model
27	logger.basicConfig(format="")	28	logger.basicConfig(format="")
28	logger.getLogger().setLevel(logger.INFO)	29	logger.getLogger().setLevel(logger.INFO)
29		30
		31	+try:
		32	+ import bitsandbytes as bnb
		33	+except ImportError:
		34	+ bnb = None
		35	+
30		36
31	def add_arguments(parser):	37	def add_arguments(parser):
32	group = parser.add_argument_group(title='Megatron saver')	38	group = parser.add_argument_group(title='Megatron saver')
	@@ -407,6 +413,21 @@ def set_model_output_layer(model_mg, msg, md, **kwargs):
407	model_mg.set_output_layer_bias(**kwargs, data=output_layer_bs[tp_rank])	413	model_mg.set_output_layer_bias(**kwargs, data=output_layer_bs[tp_rank])
408		414
409		415
		416	+def _replace_bnb_4bit_in_layer(layer):
		417	+ for name, module in layer.named_modules():
		418	+ if isinstance(module, (tpl.ColumnParallelLinear, tpl.RowParallelLinear)):
		419	+ module.weight = bnb.nn.Params4bit(
		420	+ module.weight.data,
		421	+ requires_grad=module.weight.data.requires_grad,
		422	+ quant_type="nf4"
		423	+ ).to("npu")
		424	+
		425	+
		426	+def replace_layers_parameter_to_bnb_4bit(model) -> None:
		427	+ for layer in model.decoder.layers:
		428	+ _replace_bnb_4bit_in_layer(layer)
		429	+
		430	+
410	def set_model_rm_head(model_mg, msg, md, **kwargs):	431	def set_model_rm_head(model_mg, msg, md, **kwargs):
411	margs = model_mg.get_args()	432	margs = model_mg.get_args()
412	tp_size = margs.tensor_model_parallel_size	433	tp_size = margs.tensor_model_parallel_size
	@@ -440,6 +461,8 @@ def save_model(model_mg, md, **kwargs):
440	for vp_rank in range(virtual_pipeline_model_parallel_size):	461	for vp_rank in range(virtual_pipeline_model_parallel_size):
441	kwargs["vp_rank"] = vp_rank	462	kwargs["vp_rank"] = vp_rank
442	vp_models.append(model_mg.get_model_item(**kwargs))	463	vp_models.append(model_mg.get_model_item(**kwargs))
		464	+ if args_cmd.qlora_nf4 and args_cmd.save_model_type == 'mg':
		465	+ replace_layers_parameter_to_bnb_4bit(vp_models[vp_rank])
443	if args_cmd.save_model_type == 'mg':	466	if args_cmd.save_model_type == 'mg':
444	if margs.noop_layers:	467	if margs.noop_layers:
445	for layer_idx in margs.noop_layers:	468	for layer_idx in margs.noop_layers:

Mmindspeed_llm/tasks/megatron_adaptor.py+33-2

	@@ -416,6 +416,30 @@ class CoreAdaptation(MegatronAdaptationABC):
416	# For recompute-in-advance	416	# For recompute-in-advance
417	from mindspeed.core.tensor_parallel.random import checkpoint_wrapper	417	from mindspeed.core.tensor_parallel.random import checkpoint_wrapper
418	MegatronAdaptation.register('megatron.core.tensor_parallel.random.checkpoint', checkpoint_wrapper)	418	MegatronAdaptation.register('megatron.core.tensor_parallel.random.checkpoint', checkpoint_wrapper)
		419	+ # For QLoRA
		420	+ from mindspeed_llm.tasks.posttrain.lora.utils import is_enable_qlora
		421	+ if is_enable_qlora(MegatronAdaptation.get_args()):
		422	+ from mindspeed_llm.tasks.posttrain.lora.qlora import (parallel_linear_init_wrapper,
		423	+ linear_with_frozen_weight_forward,
		424	+ linear_with_frozen_weight_backward,
		425	+ parallel_linear_save_to_state_dict_wrapper,
		426	+ parallel_linear_load_from_state_dict_wrapper)
		427	+ MegatronAdaptation.register('megatron.core.tensor_parallel.layers.ColumnParallelLinear.__init__',
		428	+ parallel_linear_init_wrapper)
		429	+ MegatronAdaptation.register('megatron.core.tensor_parallel.layers.RowParallelLinear.__init__',
		430	+ parallel_linear_init_wrapper)
		431	+ MegatronAdaptation.register('megatron.core.tensor_parallel.layers.LinearWithFrozenWeight.forward',
		432	+ linear_with_frozen_weight_forward)
		433	+ MegatronAdaptation.register('megatron.core.tensor_parallel.layers.LinearWithFrozenWeight.backward',
		434	+ linear_with_frozen_weight_backward)
		435	+ MegatronAdaptation.register('megatron.core.tensor_parallel.layers.ColumnParallelLinear._save_to_state_dict',
		436	+ parallel_linear_save_to_state_dict_wrapper)
		437	+ MegatronAdaptation.register('megatron.core.tensor_parallel.layers.RowParallelLinear._save_to_state_dict',
		438	+ parallel_linear_save_to_state_dict_wrapper)
		439	+ MegatronAdaptation.register('megatron.core.tensor_parallel.layers.ColumnParallelLinear._load_from_state_dict',
		440	+ parallel_linear_load_from_state_dict_wrapper)
		441	+ MegatronAdaptation.register('megatron.core.tensor_parallel.layers.RowParallelLinear._load_from_state_dict',
		442	+ parallel_linear_load_from_state_dict_wrapper)
419		443
420	def patch_parallel_state(self):	444	def patch_parallel_state(self):
421	import megatron	445	import megatron
	@@ -648,11 +672,18 @@ class LegacyAdaptation(MegatronAdaptationABC):
648	MegatronAdaptation.register('megatron.training.initialize.initialize_megatron', initialize_megatron)	672	MegatronAdaptation.register('megatron.training.initialize.initialize_megatron', initialize_megatron)
649		673
650	def patch_training(self):	674	def patch_training(self):
651	- from ..training import get_model_wrapper, train	675	+ from ..training import train
652	from ..training.checkpointing import load_checkpoint_wrapper	676	from ..training.checkpointing import load_checkpoint_wrapper
653	from ..legacy.data import build_pretraining_data_loader	677	from ..legacy.data import build_pretraining_data_loader
		678	+ from mindspeed_llm.tasks.posttrain.lora.utils import is_enable_qlora
		679	+
		680	+ if is_enable_qlora(MegatronAdaptation.get_args()):
		681	+ from mindspeed_llm.tasks.posttrain.lora.qlora import get_model
		682	+ MegatronAdaptation.register('megatron.training.training.get_model', get_model)
		683	+ else:
		684	+ from ..training import get_model_wrapper
		685	+ MegatronAdaptation.register('megatron.training.training.get_model', get_model_wrapper)
654		686
655	- MegatronAdaptation.register('megatron.training.training.get_model', get_model_wrapper)
656	MegatronAdaptation.register('megatron.training.training.build_pretraining_data_loader',	687	MegatronAdaptation.register('megatron.training.training.build_pretraining_data_loader',
657	build_pretraining_data_loader)	688	build_pretraining_data_loader)
658	MegatronAdaptation.register('megatron.training.training.train', train)	689	MegatronAdaptation.register('megatron.training.training.train', train)

Mmindspeed_llm/tasks/posttrain/lora/cc_lora_forward.py+33-12

	@@ -6,6 +6,17 @@ from megatron.core.parallel_state import (
6	get_tensor_model_parallel_world_size,	6	get_tensor_model_parallel_world_size,
7	)	7	)
8		8
		9	+try:
		10	+ import bitsandbytes as bnb
		11	+except ImportError:
		12	+ bnb = None
		13	+
		14	+
		15	+def dequantize(weight, dtype, device):
		16	+ if not hasattr(weight, "quant_state"):
		17	+ return weight
		18	+ return bnb.functional.dequantize_4bit(weight.data, weight.quant_state).to(device).to(dtype)
		19	+
9		20
10	def get_communication_output(input_, reduce_tensor=False):	21	def get_communication_output(input_, reduce_tensor=False):
11	tp_world_size = get_tensor_model_parallel_world_size()	22	tp_world_size = get_tensor_model_parallel_world_size()
	@@ -69,7 +80,8 @@ class _FusedColumnSeqParallelLoRAFunction(torch.autograd.Function):
69	total_input, handle = _gather_along_first_dim_async(input_)	80	total_input, handle = _gather_along_first_dim_async(input_)
70	weight_a_scale = weight_a * scaling	81	weight_a_scale = weight_a * scaling
71	ax = torch.matmul(input_, weight_a_scale.t())	82	ax = torch.matmul(input_, weight_a_scale.t())
72	- weight_combine = weight + weight_b @ weight_a_scale	83	+ weight_tmp = dequantize(weight, weight_b.dtype, weight_b.device)
		84	+ weight_combine = weight_tmp + weight_b @ weight_a_scale
73	handle.wait()	85	handle.wait()
74	output = torch.matmul(total_input, weight_combine.t())	86	output = torch.matmul(total_input, weight_combine.t())
75	ctx.save_for_backward(input_, ax, weight, weight_a_scale, weight_b)	87	ctx.save_for_backward(input_, ax, weight, weight_a_scale, weight_b)
	@@ -89,7 +101,8 @@ class _FusedColumnSeqParallelLoRAFunction(torch.autograd.Function):
89	delta_weight = weight_b @ weight_a_scale	101	delta_weight = weight_b @ weight_a_scale
90	handle.wait()	102	handle.wait()
91	grad_ax, handle = _reduce_scatter_along_first_dim_async(grad_gax)	103	grad_ax, handle = _reduce_scatter_along_first_dim_async(grad_gax)
92	- grad_input = grad_output.matmul(weight + delta_weight)	104	+ weight_tmp = dequantize(weight, delta_weight.dtype, delta_weight.device)
		105	+ grad_input = grad_output.matmul(weight_tmp + delta_weight)
93	handle.wait()	106	handle.wait()
94	grad_sub_input, handle = _reduce_scatter_along_first_dim_async(grad_input)	107	grad_sub_input, handle = _reduce_scatter_along_first_dim_async(grad_input)
95	if is_dense:	108	if is_dense:
	@@ -116,7 +129,8 @@ class _FusedRowSeqParallelLoRAFunction(torch.autograd.Function):
116	weight_a_scale = weight_a * scaling	129	weight_a_scale = weight_a * scaling
117	ax = torch.matmul(input_, weight_a_scale.t())	130	ax = torch.matmul(input_, weight_a_scale.t())
118	rax, handle = _reduce_scatter_along_first_dim_async(ax)	131	rax, handle = _reduce_scatter_along_first_dim_async(ax)
119	- weight_combine = weight + weight_b @ weight_a_scale	132	+ weight_tmp = dequantize(weight, weight_b.dtype, weight_b.device)
		133	+ weight_combine = weight_tmp + weight_b @ weight_a_scale
120	if input_.dim() == 3:	134	if input_.dim() == 3:
121	reshape = True	135	reshape = True
122	seq_len, batch, d = input_.shape[:]	136	seq_len, batch, d = input_.shape[:]
	@@ -160,8 +174,9 @@ class _FusedRowSeqParallelLoRAFunction(torch.autograd.Function):
160	input_ = input_.reshape(-1, input_.shape[-1])	174	input_ = input_.reshape(-1, input_.shape[-1])
161	else:	175	else:
162	grad_output_ = grad_output	176	grad_output_ = grad_output
		177	+ weight_tmp = dequantize(weight, grad_output_.dtype, grad_output_.device)
163	grad_input, grad_total_output = torch_npu.npu_all_gather_base_mm(	178	grad_input, grad_total_output = torch_npu.npu_all_gather_base_mm(
164	- grad_output_, weight, ctx.hcomm_info, ctx.world_size, bias=None, gather_index=0, gather_output=True	179	+ grad_output_, weight_tmp, ctx.hcomm_info, ctx.world_size, bias=None, gather_index=0, gather_output=True
165	)	180	)
166	grad_ax = grad_total_output.matmul(weight_b)	181	grad_ax = grad_total_output.matmul(weight_b)
167	grad_weight_a, grad_weight_b = lora_backward(grad_output_, input_b, grad_ax, input_, ctx.scaling)	182	grad_weight_a, grad_weight_b = lora_backward(grad_output_, input_b, grad_ax, input_, ctx.scaling)
	@@ -187,7 +202,8 @@ class _FusedRowNoSeqParallelLoRAFunction(torch.autograd.Function):
187	weight_a_scale = weight_a * scaling	202	weight_a_scale = weight_a * scaling
188	ax = torch.matmul(input_, weight_a_scale.t())	203	ax = torch.matmul(input_, weight_a_scale.t())
189	rax, handle = _reduce_async(ax)	204	rax, handle = _reduce_async(ax)
190	- output = torch.matmul(input_, weight.t())	205	+ weight_tmp = dequantize(weight, input_.dtype, input_.device)
		206	+ output = torch.matmul(input_, weight_tmp.t())
191	handle.wait()	207	handle.wait()
192	output_parallel, handle = _reduce_async(output)	208	output_parallel, handle = _reduce_async(output)
193	bx = torch.matmul(rax, weight_b.t())	209	bx = torch.matmul(rax, weight_b.t())
	@@ -208,7 +224,8 @@ class _FusedRowNoSeqParallelLoRAFunction(torch.autograd.Function):
208	grad_output_ = grad_output	224	grad_output_ = grad_output
209	grad_ax = grad_output_.matmul(weight_b)	225	grad_ax = grad_output_.matmul(weight_b)
210	grad_weight_a, grad_weight_b = lora_backward(grad_output_, input_b, grad_ax, input_, ctx.scaling)	226	grad_weight_a, grad_weight_b = lora_backward(grad_output_, input_b, grad_ax, input_, ctx.scaling)
211	- grad_input = grad_output.matmul(weight)	227	+ weight_tmp = dequantize(weight, grad_output.dtype, grad_output.device)
		228	+ grad_input = grad_output.matmul(weight_tmp)
212	grad_input += grad_ax.matmul(weight_a_scale).view_as(grad_input)	229	grad_input += grad_ax.matmul(weight_a_scale).view_as(grad_input)
213	return grad_input, None, grad_weight_a, grad_weight_b, None	230	return grad_input, None, grad_weight_a, grad_weight_b, None
214		231
	@@ -219,7 +236,8 @@ class _FusedColumnNoSeqParallelLoRAFunction(torch.autograd.Function):
219	@staticmethod	236	@staticmethod
220	def forward(ctx, input_, weight, weight_a, weight_b, scaling):	237	def forward(ctx, input_, weight, weight_a, weight_b, scaling):
221	weight_a_scale = weight_a * scaling	238	weight_a_scale = weight_a * scaling
222	- output = torch.matmul(input_, weight.t())	239	+ weight_tmp = dequantize(weight, input_.dtype, input_.device)
		240	+ output = torch.matmul(input_, weight_tmp.t())
223	ax = torch.matmul(input_, weight_a_scale.t())	241	ax = torch.matmul(input_, weight_a_scale.t())
224	bx = torch.matmul(ax, weight_b.t())	242	bx = torch.matmul(ax, weight_b.t())
225	output += bx	243	output += bx
	@@ -238,7 +256,8 @@ class _FusedColumnNoSeqParallelLoRAFunction(torch.autograd.Function):
238	grad_output_ = grad_output	256	grad_output_ = grad_output
239	grad_ax = grad_output_.matmul(weight_b)	257	grad_ax = grad_output_.matmul(weight_b)
240	grad_ax, handle = _reduce_async(grad_ax)	258	grad_ax, handle = _reduce_async(grad_ax)
241	- grad_input = grad_output.matmul(weight + weight_b @ weight_a_scale)	259	+ weight_tmp = dequantize(weight, weight_b.dtype, weight_b.device)
		260	+ grad_input = grad_output.matmul(weight_tmp + weight_b @ weight_a_scale)
242	handle.wait()	261	handle.wait()
243	grad_input, handle = _reduce_async(grad_input)	262	grad_input, handle = _reduce_async(grad_input)
244	grad_weight_a, grad_weight_b = lora_backward(grad_output_, input_b, grad_ax, input_, ctx.scaling)	263	grad_weight_a, grad_weight_b = lora_backward(grad_output_, input_b, grad_ax, input_, ctx.scaling)
	@@ -258,14 +277,15 @@ class _FusedBaseParallelLoRAFunction(torch.autograd.Function):
258	seq_size, d = input_.shape[:]	277	seq_size, d = input_.shape[:]
259	weight_a_scale = weight_a * scaling	278	weight_a_scale = weight_a * scaling
260	ax = torch.matmul(input_, weight_a_scale.t())	279	ax = torch.matmul(input_, weight_a_scale.t())
		280	+ weight_tmp = dequantize(weight, input_.dtype, input_.device)
261	if seq_size < d:	281	if seq_size < d:
262	ctx.combine = False	282	ctx.combine = False
263	- output = torch.matmul(input_, weight.t())	283	+ output = torch.matmul(input_, weight_tmp.t())
264	bx = torch.matmul(ax, weight_b.t())	284	bx = torch.matmul(ax, weight_b.t())
265	output += bx	285	output += bx
266	else:	286	else:
267	ctx.combine = True	287	ctx.combine = True
268	- weight_combine = weight + weight_b @ weight_a_scale	288	+ weight_combine = weight_tmp + weight_b @ weight_a_scale
269	output = torch.matmul(input_, weight_combine.t())	289	output = torch.matmul(input_, weight_combine.t())
270	ctx.save_for_backward(input_, ax, weight_a_scale, weight_b, weight)	290	ctx.save_for_backward(input_, ax, weight_a_scale, weight_b, weight)
271	ctx.scaling = scaling	291	ctx.scaling = scaling
	@@ -282,10 +302,11 @@ class _FusedBaseParallelLoRAFunction(torch.autograd.Function):
282	grad_output_ = grad_output	302	grad_output_ = grad_output
283	grad_ax = grad_output_.matmul(weight_b)	303	grad_ax = grad_output_.matmul(weight_b)
284	grad_weight_a, grad_weight_b = lora_backward(grad_output_, input_b, grad_ax, input_, ctx.scaling)	304	grad_weight_a, grad_weight_b = lora_backward(grad_output_, input_b, grad_ax, input_, ctx.scaling)
		305	+ weight_tmp = dequantize(weight, grad_output.dtype, grad_output.device)
285	if ctx.combine:	306	if ctx.combine:
286	- grad_input = grad_output.matmul((weight + weight_b @ weight_a_scale))	307	+ grad_input = grad_output.matmul((weight_tmp + weight_b @ weight_a_scale))
287	else:	308	else:
288	- grad_input = grad_output.matmul(weight)	309	+ grad_input = grad_output.matmul(weight_tmp)
289	grad_input += grad_ax.matmul(weight_a_scale).view_as(grad_input)	310	grad_input += grad_ax.matmul(weight_a_scale).view_as(grad_input)
290	return grad_input, None, grad_weight_a, grad_weight_b, None	311	return grad_input, None, grad_weight_a, grad_weight_b, None
291		312

Amindspeed_llm/tasks/posttrain/lora/qlora.py+209-0

@@ -0,0 +1,209 @@
	1	+# Copyright (c) 2025, HUAWEI CORPORATION. All rights reserved.
	2	+import torch
	3	+import torch_npu
	4	+
	5	+from megatron.training import get_args
	6	+from megatron.core import mpu
	7	+from megatron.core.utils import get_model_config
	8	+from megatron.core.enums import ModelType
	9	+from megatron.core.distributed import DistributedDataParallel as DDP
	10	+from megatron.core.parallel_state import get_tensor_model_parallel_group
	11	+from mindspeed_llm.training.training import model_provider_func_wrapper
	12	+from mindspeed_llm.tasks.posttrain.lora.utils import is_enable_qlora
	13	+
	14	+try:
	15	+ import bitsandbytes as bnb
	16	+except ImportError:
	17	+ bnb = None
	18	+
	19	+
	20	+def parallel_linear_init_wrapper(fn):
	21	+ def wrapper(self, input_size, output_size, **kwargs):
	22	+ fn(self, input_size, output_size, **kwargs)
	23	+ if is_enable_qlora():
	24	+ self.weight.data = self.weight.data.to("cpu")
	25	+ return wrapper
	26	+
	27	+
	28	+def linear_with_frozen_weight_forward(
	29	+ ctx, input_, weight, bias, allreduce_dgrad
	30	+ ):
	31	+ ctx.save_for_backward(weight)
	32	+ ctx.allreduce_dgrad = allreduce_dgrad
	33	+ if hasattr(weight, "quant_state"):
	34	+ weight_tmp = bnb.functional.dequantize_4bit(weight.data, weight.quant_state).to(input_.dtype)
	35	+ else:
	36	+ weight_tmp = weight
	37	+ output = torch.matmul(input_, weight_tmp.t())
	38	+ if bias is not None:
	39	+ output = output + bias
	40	+ return output
	41	+
	42	+
	43	+def linear_with_frozen_weight_backward(ctx, grad_output):
	44	+ (weight,) = ctx.saved_tensors
	45	+ if hasattr(weight, "quant_state"):
	46	+ weight_tmp = bnb.functional.dequantize_4bit(weight.data, weight.quant_state).to(grad_output.dtype)
	47	+ else:
	48	+ weight_tmp = weight
	49	+ grad_input = grad_output.matmul(weight_tmp)
	50	+ if ctx.allreduce_dgrad:
	51	+ # All-reduce. Note: here async and sync are effectively the same.
	52	+ torch.distributed.all_reduce(grad_input, group=get_tensor_model_parallel_group())
	53	+
	54	+ return grad_input, None, None, None
	55	+
	56	+
	57	+def parallel_linear_save_to_state_dict_wrapper(fn):
	58	+ def wrapper(self, destination, prefix, keep_vars):
	59	+ """
	60	+ save weight and bias,
	61	+ then fill state_dict with components of quant_state
	62	+ """
	63	+ fn(self, destination, prefix, keep_vars)
	64	+ if getattr(self.weight, "quant_state", None) is not None:
	65	+ for k, v in self.weight.quant_state.as_dict(packed=True).items():
	66	+ destination[prefix + "weight." + k] = v if keep_vars else v.detach()
	67	+
	68	+ return wrapper
	69	+
	70	+
	71	+def parallel_linear_load_from_state_dict_wrapper(fn):
	72	+ def wrapper(self, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs):
	73	+ if any(['bitsandbytes' in i for i in state_dict.keys()]): # is quantized linear
	74	+ qs_dict = {key: v for k, v in state_dict.items() if (key := k.replace(prefix, "")) != '_extra_state'}
	75	+ self.weight = bnb.nn.Params4bit.from_prequantized(
	76	+ data=qs_dict.get('weight'),
	77	+ quantized_stats={key.replace('weight.', ''): qs_dict[key] for key in qs_dict if key != 'weight'},
	78	+ requires_grad=False,
	79	+ device='npu')
	80	+ fn(self, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs)
	81	+ return wrapper
	82	+
	83	+
	84	+def get_model(model_provider_func, model_type=ModelType.encoder_or_decoder, wrap_with_ddp=True):
	85	+ """Build the model."""
	86	+ from megatron.core import tensor_parallel
	87	+ from megatron.legacy.model import Float16Module
	88	+ from megatron.core.distributed import DistributedDataParallelConfig
	89	+
	90	+ tpl = tensor_parallel.layers
	91	+ model_provider_func = model_provider_func_wrapper(model_provider_func)
	92	+ args = get_args()
	93	+ args.model_type = model_type
	94	+
	95	+ # Build model.
	96	+ if mpu.get_pipeline_model_parallel_world_size() > 1 and \
	97	+ args.virtual_pipeline_model_parallel_size is not None:
	98	+ assert model_type != ModelType.encoder_and_decoder, \
	99	+ "Interleaved schedule not supported for model with both encoder and decoder"
	100	+ model = []
	101	+ for i in range(args.virtual_pipeline_model_parallel_size):
	102	+ mpu.set_virtual_pipeline_model_parallel_rank(i)
	103	+ # Set pre_process and post_process only after virtual rank is set.
	104	+ pre_process = mpu.is_pipeline_first_stage()
	105	+ post_process = mpu.is_pipeline_last_stage()
	106	+ this_model = model_provider_func(
	107	+ pre_process=pre_process,
	108	+ post_process=post_process
	109	+ )
	110	+ this_model.model_type = model_type
	111	+ model.append(this_model)
	112	+ else:
	113	+ pre_process = mpu.is_pipeline_first_stage()
	114	+ post_process = mpu.is_pipeline_last_stage()
	115	+ add_encoder = True
	116	+ add_decoder = True
	117	+ if model_type == ModelType.encoder_and_decoder:
	118	+ if mpu.get_pipeline_model_parallel_world_size() > 1:
	119	+ assert args.pipeline_model_parallel_split_rank is not None, \
	120	+ "Split rank needs to be specified for model with both encoder and decoder"
	121	+ rank = mpu.get_pipeline_model_parallel_rank()
	122	+ split_rank = args.pipeline_model_parallel_split_rank
	123	+ world_size = mpu.get_pipeline_model_parallel_world_size()
	124	+ pre_process = rank == 0 or rank == split_rank
	125	+ post_process = (rank == (split_rank - 1)) or (
	126	+ rank == (world_size - 1))
	127	+ add_encoder = mpu.is_pipeline_stage_before_split()
	128	+ add_decoder = mpu.is_pipeline_stage_after_split()
	129	+ model = model_provider_func(
	130	+ pre_process=pre_process,
	131	+ post_process=post_process,
	132	+ add_encoder=add_encoder,
	133	+ add_decoder=add_decoder)
	134	+ else:
	135	+ model = model_provider_func(
	136	+ pre_process=pre_process,
	137	+ post_process=post_process
	138	+ )
	139	+ model.model_type = model_type
	140	+
	141	+ if not isinstance(model, list):
	142	+ model = [model]
	143	+
	144	+ # Set tensor model parallel attributes if not set.
	145	+ # Only parameters that are already tensor model parallel have these
	146	+ # attributes set for them. We should make sure the default attributes
	147	+ # are set for all params so the optimizer can use them.
	148	+ for model_module in model:
	149	+ for param in model_module.parameters():
	150	+ tensor_parallel.set_defaults_if_not_set_tensor_model_parallel_attributes(param)
	151	+
	152	+ # Print number of parameters.
	153	+ if mpu.get_data_parallel_rank() == 0:
	154	+ print(' > number of parameters on (tensor, pipeline) '
	155	+ 'model parallel rank ({}, {}): {}'.format(
	156	+ mpu.get_tensor_model_parallel_rank(),
	157	+ mpu.get_pipeline_model_parallel_rank(),
	158	+ sum([sum([p.nelement() for p in model_module.parameters()])
	159	+ for model_module in model])), flush=True)
	160	+
	161	+ # start of megatron_adaptation,
	162	+ # here we keep the main model's linear layers on CPU to avoid OOM in QLoRA.
	163	+ # GPU allocation.
	164	+ for model_module in model:
	165	+ if is_enable_qlora():
	166	+ for name, module in model_module.base_model.named_modules():
	167	+ if not hasattr(module, "weight") or hasattr(module, "base_layer"):
	168	+ continue
	169	+
	170	+ is_lora_adapter = any(substring in name for substring in ["lora_A", "lora_B"])
	171	+ is_target_linear = (
	172	+ isinstance(module, (tpl.ColumnParallelLinear, tpl.RowParallelLinear, torch.nn.Linear))
	173	+ and "layers" in name
	174	+ )
	175	+
	176	+ if not (is_target_linear and not is_lora_adapter):
	177	+ module.weight.data = module.weight.data.to(torch.cuda.current_device())
	178	+ else:
	179	+ model_module.cuda(torch.cuda.current_device())
	180	+ # end of megatron_adaptation
	181	+
	182	+ # Fp16 conversion.
	183	+ if args.fp16 or args.bf16:
	184	+ model = [Float16Module(model_module, args) for model_module in model]
	185	+
	186	+ if wrap_with_ddp:
	187	+ config = get_model_config(model[0])
	188	+ ddp_config = DistributedDataParallelConfig(
	189	+ grad_reduce_in_fp32=args.accumulate_allreduce_grads_in_fp32,
	190	+ overlap_grad_reduce=args.overlap_grad_reduce,
	191	+ use_distributed_optimizer=args.use_distributed_optimizer,
	192	+ check_for_nan_in_grad=args.check_for_nan_in_loss_and_grad,
	193	+ bucket_size=args.ddp_bucket_size)
	194	+ model = [DDP(config,
	195	+ ddp_config,
	196	+ model_chunk,
	197	+ data_parallel_group=mpu.get_data_parallel_group(with_context_parallel=True),
	198	+ expert_data_parallel_group=mpu.get_data_modulo_expert_parallel_group(),
	199	+ # Turn off bucketing for model_chunk 2 onwards, since communication for these
	200	+ # model chunks is overlapped with compute anyway.
	201	+ disable_bucketing=(model_chunk_idx > 0))
	202	+ for (model_chunk_idx, model_chunk) in enumerate(model)]
	203	+
	204	+ # Broadcast params from data parallel src rank to other data parallel ranks.
	205	+ if args.data_parallel_random_init:
	206	+ for model_module in model:
	207	+ model_module.broadcast_params()
	208	+
	209	+ return model

Mmindspeed_llm/tasks/posttrain/lora/utils.py+7-0

	@@ -28,6 +28,13 @@ def is_enable_lora():
28	return False	28	return False
29		29
30		30
		31	+def is_enable_qlora(args=None):
		32	+ args = args if args else get_args()
		33	+ if hasattr(args, 'qlora') and args.qlora:
		34	+ return True
		35	+ return False
		36	+
		37	+
31	def merge_dicts(dict1, dict2):	38	def merge_dicts(dict1, dict2):
32	result = dict1	39	result = dict1
33	for key, value in dict2.items():	40	for key, value in dict2.items():

Mmindspeed_llm/training/arguments.py+2-0

	@@ -302,6 +302,8 @@ def _add_lora_args(parser):
302	help='Lora register forward hook.')	302	help='Lora register forward hook.')
303	group.add_argument('--lora-fusion', action='store_true',	303	group.add_argument('--lora-fusion', action='store_true',
304	help='use fusion to accelerate lora.')	304	help='use fusion to accelerate lora.')
		305	+ group.add_argument('--qlora', action='store_true', default=False,
		306	+ help='Enable QLoRA for fine-tuning with reduced memory usage.')
305	return parser	307	return parser
306		308
307		309

Mrequirements.txt+2-1

	@@ -16,4 +16,5 @@ tiktoken
16	ray==2.10.0	16	ray==2.10.0
17	tensordict	17	tensordict
18	hydra-core	18	hydra-core
19	-codetiming	19	+codetiming
		20	+bitsandbytes-npu-beta==0.45.2

Asources/images/qlora/qlora.png+0-0

Mtests/README.md+10-2

	@@ -13,7 +13,7 @@
13	<th>Mem.</th>	13	<th>Mem.</th>
14	</tr>	14	</tr>
15	<tr>	15	<tr>
16	- <td rowspan="23">ST</td>	16	+ <td rowspan="24">ST</td>
17	<td rowspan="13">Pretrain</td>	17	<td rowspan="13">Pretrain</td>
18	<td>Mcore</td>	18	<td>Mcore</td>
19	<td>TP，PP，VPP，distributed_optimizer，o2_gradient，o2_optimizer，重计算，enable_recompute_layers_per_pp_rank，FA_TND，use_fused_rotary_pos_emb_new</td>	19	<td>TP，PP，VPP，distributed_optimizer，o2_gradient，o2_optimizer，重计算，enable_recompute_layers_per_pp_rank，FA_TND，use_fused_rotary_pos_emb_new</td>
	@@ -119,7 +119,7 @@
119	<td>Y</td>	119	<td>Y</td>
120	</tr>	120	</tr>
121	<tr>	121	<tr>
122	- <td rowspan="1">LoRA</td>	122	+ <td rowspan="2">LoRA</td>
123	<td rowspan="1">Legacy</td>	123	<td rowspan="1">Legacy</td>
124	<td>CCLoRA, TP, PP, 全重计算</td>	124	<td>CCLoRA, TP, PP, 全重计算</td>
125	<td><a href="st/shell_scripts/tune_llama2_tp2_pp4_lora_ptd.sh">tune_llama2_tp2_pp4_lora_ptd.sh</a></td>	125	<td><a href="st/shell_scripts/tune_llama2_tp2_pp4_lora_ptd.sh">tune_llama2_tp2_pp4_lora_ptd.sh</a></td>
	@@ -127,6 +127,14 @@
127	<td>Y</td>	127	<td>Y</td>
128	<td>Y</td>	128	<td>Y</td>
129	</tr>	129	</tr>
		130	+ <tr>
		131	+ <td rowspan="1">Mcore</td>
		132	+ <td>CCLoRA, QLoRA</td>
		133	+ <td><a href="st/shell_scripts/tune_llama2_tp1_pp1_qlora_ptd.sh">tune_llama2_tp1_pp1_qlora_ptd.sh</a></td>
		134	+ <td>Y</td>
		135	+ <td>Y</td>
		136	+ <td>Y</td>
		137	+ </tr>
130	<tr>	138	<tr>
131	<td rowspan="3">DPO</td>	139	<td rowspan="3">DPO</td>
132	<td rowspan="3">Mcore</td>	140	<td rowspan="3">Mcore</td>

Atests/st/baseline_results/tune_llama2_tp1_pp1_qlora_ptd.json+43-0

@@ -0,0 +1,43 @@
	1	+{
	2	+ "lm loss": [
	3	+ 1.229878,
	4	+ 1.294799,
	5	+ 1.300336,
	6	+ 1.314004,
	7	+ 1.401411,
	8	+ 1.257517,
	9	+ 1.360531,
	10	+ 1.300218,
	11	+ 1.447775,
	12	+ 1.302323,
	13	+ 1.293658,
	14	+ 1.27972,
	15	+ 1.285732,
	16	+ 1.241777,
	17	+ 1.258219
	18	+ ],
	19	+ "throughput": [
	20	+ 1475.2,
	21	+ 2820.7,
	22	+ 2957.6,
	23	+ 2819.5,
	24	+ 2983.0,
	25	+ 2860.5,
	26	+ 2890.7,
	27	+ 2989.2,
	28	+ 2802.4,
	29	+ 2824.5,
	30	+ 2904.8,
	31	+ 2803.8,
	32	+ 2887.0,
	33	+ 2872.8,
	34	+ 2987.6
	35	+ ],
	36	+ "memo info": [
	37	+ {
	38	+ "rank": 0,
	39	+ "allocated memory": 4360.6865234375,
	40	+ "max allocated memory": 10336.7099609375
	41	+ }
	42	+ ]
	43	+}

Atests/st/shell_scripts/tune_llama2_tp1_pp1_qlora_ptd.sh+112-0

@@ -0,0 +1,112 @@
	1	+#!/bin/bash
	2	+
	3	+export CUDA_DEVICE_MAX_CONNECTIONS=1
	4	+export PYTORCH_NPU_ALLOC_CONF=expandable_segments:True
	5	+export HCCL_CONNECT_TIMEOUT=1200
	6	+export HCCL_EXEC_TIMEOUT=1200
	7	+
	8	+GPUS_PER_NODE=1
	9	+MASTER_ADDR=localhost
	10	+MASTER_PORT=6080
	11	+NNODES=1
	12	+NODE_RANK=0
	13	+WORLD_SIZE=$(($GPUS_PER_NODE*$NNODES))
	14	+
	15	+basepath=$(cd `dirname $0`; cd ../../../; pwd)
	16	+
	17	+CKPT_SAVE_DIR="/data/ckpt"
	18	+CKPT_LOAD_DIR="/data/llama2-7b-tp1-pp1-nf4"
	19	+DATA_PATH="/data/tune_dataset/alpaca"
	20	+TOKENIZER_MODEL="/data/llama-2-7b-hf/"
	21	+
	22	+TP=1
	23	+PP=1
	24	+
	25	+DISTRIBUTED_ARGS=(
	26	+ --nproc_per_node $GPUS_PER_NODE
	27	+ --nnodes $NNODES
	28	+ --node_rank $NODE_RANK
	29	+ --master_addr $MASTER_ADDR
	30	+ --master_port $MASTER_PORT
	31	+)
	32	+
	33	+DIST_ALGO=(
	34	+ --tensor-model-parallel-size ${TP}
	35	+ --pipeline-model-parallel-size ${PP}
	36	+)
	37	+
	38	+MODEL_ARGS=(
	39	+ --num-layers 32
	40	+ --hidden-size 4096
	41	+ --ffn-hidden-size 11008
	42	+ --num-attention-heads 32
	43	+ --seq-length 4096
	44	+ --max-position-embeddings 4096
	45	+)
	46	+
	47	+TRAINING_ARGS=(
	48	+ --use-mcore-models
	49	+ --tokenizer-type PretrainedFromHF
	50	+ --tokenizer-name-or-path ${TOKENIZER_MODEL}
	51	+ --micro-batch-size 4
	52	+ --global-batch-size 32
	53	+ --make-vocab-size-divisible-by 1
	54	+ --padded-vocab-size 32000
	55	+ --lr 1.25e-6
	56	+ --train-iters 15
	57	+ --lr-decay-style cosine
	58	+ --untie-embeddings-and-output-weights
	59	+ --disable-bias-linear
	60	+ --attention-dropout 0.0
	61	+ --init-method-std 0.01
	62	+ --hidden-dropout 0.0
	63	+ --position-embedding-type rope
	64	+ --normalization RMSNorm
	65	+ --use-fused-rmsnorm
	66	+ --swiglu
	67	+ --use-flash-attn
	68	+ --no-masked-softmax-fusion
	69	+ --attention-softmax-in-fp32
	70	+ --min-lr 1.25e-7
	71	+ --weight-decay 1e-1
	72	+ --lr-warmup-fraction 0.01
	73	+ --clip-grad 1.0
	74	+ --adam-beta1 0.9
	75	+ --initial-loss-scale 65536
	76	+ --adam-beta2 0.95
	77	+ --no-gradient-accumulation-fusion
	78	+ --no-load-optim
	79	+ --no-load-rng
	80	+ --finetune
	81	+ --stage sft
	82	+ --is-instruction-dataset
	83	+ --lora-r 16
	84	+ --lora-alpha 32
	85	+ --lora-fusion
	86	+ --lora-target-modules linear_qkv linear_proj linear_fc1 linear_fc2
	87	+ --qlora
	88	+ --variable-seq-lengths
	89	+ --bf16
	90	+)
	91	+
	92	+DATA_ARGS=(
	93	+ --data-path $DATA_PATH
	94	+ --split 949,50,1
	95	+)
	96	+
	97	+OUTPUT_ARGS=(
	98	+ --log-interval 1
	99	+ --eval-interval 1000
	100	+ --eval-iters 0
	101	+)
	102	+
	103	+torchrun ${DISTRIBUTED_ARGS[@]} $basepath/posttrain_gpt.py \
	104	+ ${DIST_ALGO[@]} \
	105	+ ${MODEL_ARGS[@]} \
	106	+ ${TRAINING_ARGS[@]} \
	107	+ ${DATA_ARGS[@]} \
	108	+ ${OUTPUT_ARGS[@]} \
	109	+ --load ${CKPT_LOAD_DIR} \
	110	+ --finetune \
	111	+ --log-throughput \
	112	+ --distributed-backend nccl