feat: support configurable LoRA weight saving-MindSpeed-MM-AtomGit

Mdocs/en/features/lora_finetune_fsdp2.md+13-11

	@@ -50,6 +50,7 @@ training:
50	- "model.language_model.layers.{*}.mlp.up_proj"	50	- "model.language_model.layers.{*}.mlp.up_proj"
51	- "model.language_model.layers.{*}.mlp.down_proj"	51	- "model.language_model.layers.{*}.mlp.down_proj"
52	dropout: 0.0	52	dropout: 0.0
		53	+ lora_save_only: false
53	init_lora_weights: true	54	init_lora_weights: true
54	pretrained_lora_path: null	55	pretrained_lora_path: null
55	```	56	```
	@@ -63,6 +64,7 @@ training:
63	\| `alpha` \| int \| `16` \| Controls the influence ratio of LoRA weights on the original weights; a higher value means greater influence. Generally keep `α/r` at 2 \|	64	\| `alpha` \| int \| `16` \| Controls the influence ratio of LoRA weights on the original weights; a higher value means greater influence. Generally keep `α/r` at 2 \|
64	\| `target_modules` \| List[str] \| `["q_proj", "k_proj", "v_proj"]` \| Module names or wildcard patterns to which LoRA is added \|	65	\| `target_modules` \| List[str] \| `["q_proj", "k_proj", "v_proj"]` \| Module names or wildcard patterns to which LoRA is added \|
65	\| `dropout` \| float \| `0.0` \| Dropout ratio of the LoRA layer, in the range `[0, 1)` \|	66	\| `dropout` \| float \| `0.0` \| Dropout ratio of the LoRA layer, in the range `[0, 1)` \|
		67	+\| `lora_save_only` \| bool \| `false` \| When `true`, export only standalone LoRA safetensors and the adapter configuration file. When `false`, save the full model weights according to the save format: DCP saves unmerged base weights and LoRA weights, while HF saves the base weights merged with the LoRA weights. \|
66	\| `init_lora_weights` \| bool \\| str \| `True` \| Weight initialization method. `True`; `False`; or one of the following string values: `"gaussian"`, `"eva"`, `"olora"`, `"pissa"`, `"pissa_niter_[number of iters]"`, `"corda"`, `"loftq"`, `"orthogonal"` \|	68	\| `init_lora_weights` \| bool \\| str \| `True` \| Weight initialization method. `True`; `False`; or one of the following string values: `"gaussian"`, `"eva"`, `"olora"`, `"pissa"`, `"pissa_niter_[number of iters]"`, `"corda"`, `"loftq"`, `"orthogonal"` \|
67	\| `pretrained_lora_path` \| str \| `null` \| Path to pretrained LoRA weights (optional), supporting `.safetensors` and `.pt/.bin` formats \|	69	\| `pretrained_lora_path` \| str \| `null` \| Path to pretrained LoRA weights (optional), supporting `.safetensors` and `.pt/.bin` formats \|
68		70
	@@ -115,20 +117,18 @@ To load pretrained LoRA weights for resumable training, configure the `pretraine
115	training:	117	training:
116	lora:	118	lora:
117	enable: true	119	enable: true
118	- pretrained_lora_path: ./save_path/iter_xxx # Replace with the LoRA weight save path.	120	+ pretrained_lora_path: ./save_path/iter_00000xx/lora_adapter_iteration_xx.safetensors
119	```	121	```
120		122
121	## Weight Saving	123	## Weight Saving
122		124
123	-### Saving Only LoRA Weights	125	+Use `training.lora.lora_save_only` to select what to save. The default is `false`. When weight saving is enabled, the following three modes are available:
124		126
125	-During training, only the LoRA adapter weights are saved in the safetensors format. The saved file structure is as follows:	127	+\| Save Mode \| Key Configuration \| Saved Content \| Use Case \|
126	-	128	+\| :--- \| :--- \| :--- \| :--- \|
127	-```bash	129	+\| Save only LoRA \| `lora_save_only: true` \| Save LoRA weights in safetensors format and `adapter_config.json` in the weight output directory, without base weights or training state such as optimizer and RNG state \| Distributing or reusing adapters, and merging them with the base model offline \|
128	-save_path/	130	+\| Save a DCP checkpoint \| `lora_save_only: false`, `save_format: dcp` \| Save unmerged base weights and LoRA weights; optimizer, RNG, and other state are controlled by the existing save parameters \| Resuming training from a checkpoint \|
129	-├── lora_adapter.safetensors	131	+\| Save an HF checkpoint \| `lora_save_only: false`, `save_format: hf` \| Merge weights using `base + (alpha / rank) * B @ A` and save standard HF weights, without retaining separate LoRA parameters \| Inference, deployment, or publishing a full model \|
130	-└── ...
131	-```
132		132
133	## Starting Training	133	## Starting Training
134		134
	@@ -142,17 +142,19 @@ Once training starts, a LoRA configuration summary is automatically printed, inc
142		142
143	## Merging LoRA Weights into Hugging Face Weights	143	## Merging LoRA Weights into Hugging Face Weights
144		144
		145	+When `lora_save_only=true`, only LoRA weights are saved. Use the following script to merge them with the original HF weights offline:
		146	+
145	```bash	147	```bash
146	cd checkpoint/common	148	cd checkpoint/common
147	python merge_lora_safetensors_to_base.py \	149	python merge_lora_safetensors_to_base.py \
148	--base_hf_dir ./Qwen3.5-27B \	150	--base_hf_dir ./Qwen3.5-27B \
149	- --lora_safetensors ./save_path/lora_adapter_iteration_10.safetensors \	151	+ --lora_safetensors ./save_path/iter_00000xx/lora_adapter_iteration_xx.safetensors \
150	--save_merged_hf_dir ./merged_qwen3_5_27B_lora	152	--save_merged_hf_dir ./merged_qwen3_5_27B_lora
151	```	153	```
152		154
153	## Resuming Training from Checkpoint for LoRA	155	## Resuming Training from Checkpoint for LoRA
154		156
155	-When resuming training from checkpoint, the `load` path in the YAML configuration file must point to the checkpoint path saved by the previous training run. The previous training run must have `no_save_optim` and `no_save_rng` set to `false.` To resume training, set `no_load_optim` and `no_load_rng` to `false` to restore the optimizer state. After resume training is complete, use the weight conversion script to merge the LoRA weights into Hugging Face weights.	157	+To fully resume training, the previous checkpoint must have been saved with `save_format=dcp`, `lora.lora_save_only=false`, and both `no_save_optim` and `no_save_rng` set to `false`. When resuming, set `load` to the checkpoint root directory containing `latest_checkpointed_iteration.txt`, and set both `no_load_optim` and `no_load_rng` to `false`. DCP checkpoints contain the full base weights and LoRA weights, so `pretrained_lora_path` does not need to be set.
156		158
157	## Notes	159	## Notes
158		160

Mdocs/zh/features/lora_finetune_fsdp2.md+13-11

	@@ -44,6 +44,7 @@ training:
44	# 配合 model.freeze 排除组件（见下文 target_modules 配置说明）	44	# 配合 model.freeze 排除组件（见下文 target_modules 配置说明）
45	target_modules: all-linear	45	target_modules: all-linear
46	dropout: 0.0	46	dropout: 0.0
		47	+ lora_save_only: false
47	init_lora_weights: true	48	init_lora_weights: true
48	pretrained_lora_path: null	49	pretrained_lora_path: null
49	```	50	```
	@@ -57,6 +58,7 @@ training:
57	\| `alpha` \| int \| `16` \| 控制 LoRA 权重对原始权重的影响比例，数值越高影响越大。一般保持 `α/r` 为 2 \|	58	\| `alpha` \| int \| `16` \| 控制 LoRA 权重对原始权重的影响比例，数值越高影响越大。一般保持 `α/r` 为 2 \|
58	\| `target_modules` \| str \\| List[str] \| `["q_proj", "k_proj", "v_proj"]` \| 需要添加 LoRA 的模块名称，或者通配符模式，或特殊关键字 `all-linear` \|	59	\| `target_modules` \| str \\| List[str] \| `["q_proj", "k_proj", "v_proj"]` \| 需要添加 LoRA 的模块名称，或者通配符模式，或特殊关键字 `all-linear` \|
59	\| `dropout` \| float \| `0.0` \| LoRA 层的 dropout 比例，取值范围 `[0, 1)` \|	60	\| `dropout` \| float \| `0.0` \| LoRA 层的 dropout 比例，取值范围 `[0, 1)` \|
		61	+\| `lora_save_only` \| bool \| `false` \| `true` 时只导出独立 LoRA safetensors 和 adapter 配置文件；`false` 时根据保存类型存储全量模型权重，DCP 格式时保存基础权重和 LoRA 权重，HF 格式时保存基础权重与 LoRA 权重融合后的权重。 \|
60	\| `init_lora_weights` \| bool \\| str \| `True` \| 权重初始化方式。`True`；`False`；或选择以下字符串值：`"gaussian"`, `"eva"`, `"olora"`, `"pissa"`, `"pissa_niter_[number of iters]"`, `"corda"`, `"loftq"`, `"orthogonal"` \|	62	\| `init_lora_weights` \| bool \\| str \| `True` \| 权重初始化方式。`True`；`False`；或选择以下字符串值：`"gaussian"`, `"eva"`, `"olora"`, `"pissa"`, `"pissa_niter_[number of iters]"`, `"corda"`, `"loftq"`, `"orthogonal"` \|
61	\| `pretrained_lora_path` \| str \| `null` \| 预训练 LoRA 权重路径（可选），支持 `.safetensors` 和 `.pt/.bin` 格式 \|	63	\| `pretrained_lora_path` \| str \| `null` \| 预训练 LoRA 权重路径（可选），支持 `.safetensors` 和 `.pt/.bin` 格式 \|
62	\| `disable_peft_moe_conversion` \| bool \| `true` \| 屏蔽 PEFT 对 MoE 模型 `gate_proj`/`up_proj`/`down_proj` 的 `target_modules→target_parameters` 自动重定向，使 LoRA 打在 `shared_expert` 的 `nn.Linear` 而非路由专家的 `nn.Parameter` 上。仅 MoE 模型相关 \|	64	\| `disable_peft_moe_conversion` \| bool \| `true` \| 屏蔽 PEFT 对 MoE 模型 `gate_proj`/`up_proj`/`down_proj` 的 `target_modules→target_parameters` 自动重定向，使 LoRA 打在 `shared_expert` 的 `nn.Linear` 而非路由专家的 `nn.Parameter` 上。仅 MoE 模型相关 \|
	@@ -128,20 +130,18 @@ target_modules:
128	training:	130	training:
129	lora:	131	lora:
130	enable: true	132	enable: true
131	- pretrained_lora_path: ./save_path/iter_xxx # 替换为 LoRA 权重保存路径	133	+ pretrained_lora_path: ./save_path/iter_00000xx/lora_adapter_iteration_xx.safetensors
132	```	134	```
133		135
134	## 权重保存	136	## 权重保存
135		137
136	-### 仅保存 LoRA 权重	138	+通过 `training.lora.lora_save_only` 选择保存范围，默认值为 `false`，开启权重保存后有以下三种保存方式：
137		139
138	-训练过程中仅保存 LoRA 适配器权重，保存格式为 safetensors，保存的文件结构：	140	+\| 保存方式 \| 关键配置 \| 保存内容 \| 适用场景 \|
139	-	141	+\| :--- \| :--- \| :--- \| :--- \|
140	-```bash	142	+\| 仅保存 LoRA \| `lora_save_only: true` \| 在权重保存目录下保存 LoRA 权重 safetensors 和 `adapter_config.json`，不保存基础权重及 optimizer、RNG 等训练状态 \| 分发或复用 adapter，以及后续与基础模型离线融合 \|
141	-save_path/	143	+\| 保存 DCP checkpoint \| `lora_save_only: false`、`save_format: dcp` \| 保存未融合的基础权重和 LoRA 权重；optimizer、RNG 等状态由原有保存参数控制 \| 断点续训 \|
142	-├── lora_adapter.safetensors	144	+\| 保存 HF checkpoint \| `lora_save_only: false`、`save_format: hf` \| 按 `base + (alpha / rank) * B @ A` 融合并保存标准 HF 权重，不保留独立 LoRA 参数 \| 推理、部署或发布完整模型 \|
143	-└── ...
144	-```
145		145
146	## 启动训练	146	## 启动训练
147		147
	@@ -155,17 +155,19 @@ bash examples/qwen3_5/finetune_qwen3_5_xxB.sh
155		155
156	## 合并lora权重到HuggingFace权重	156	## 合并lora权重到HuggingFace权重
157		157
		158	+开启 `lora_save_only=true` 参数后仅保存 LoRA 权重，可通过以下脚本将其与原始 HF 权重离线合并：
		159	+
158	```bash	160	```bash
159	cd checkpoint/common	161	cd checkpoint/common
160	python merge_lora_safetensors_to_base.py \	162	python merge_lora_safetensors_to_base.py \
161	--base_hf_dir ./Qwen3.5-27B \	163	--base_hf_dir ./Qwen3.5-27B \
162	- --lora_safetensors ./save_path/lora_adapter_iteration_10.safetensors \	164	+ --lora_safetensors ./save_path/iter_00000xx/lora_adapter_iteration_xx.safetensors \
163	--save_merged_hf_dir ./merged_qwen3_5_27B_lora	165	--save_merged_hf_dir ./merged_qwen3_5_27B_lora
164	```	166	```
165		167
166	## lora断点续训	168	## lora断点续训
167		169
168	-断点续训时，yaml配置文件中`load`路径需要指向上次训练保存的 checkpoint 路径。断点续训前一次的训练必须配置`no_save_optim`、`no_save_rng`为false，断点续训时`no_load_optim`、`no_load_rng`设置为false，才能恢复优化器状态。断点续训完成后可使用权重转换脚本，合并lora权重到HuggingFace权重。	170	+完整断点续训需在上次保存时设置 `save_format=dcp`、`lora.lora_save_only=false`，并将 `no_save_optim`、`no_save_rng` 设为 `false`。恢复时，`load` 指向包含 `latest_checkpointed_iteration.txt` 的 checkpoint 根目录，`no_load_optim`、`no_load_rng` 设为 `false`。DCP 权重包含完整基础权重和 LoRA 权重，无需再设置 `pretrained_lora_path`。
169		171
170	## lora微调支持模型	172	## lora微调支持模型
171		173

Mmindspeed_mm/fsdp/checkpoint/convert.py+10-5

	@@ -274,6 +274,11 @@ class Qwen35WeightTransformPipeline(WeightTransformPipeline):
274	tensor = permute_moe_expert(key, tensor, self.expert_weight_patterns)	274	tensor = permute_moe_expert(key, tensor, self.expert_weight_patterns)
275	return key, tensor	275	return key, tensor
276		276
		277	+ @staticmethod
		278	+ def _to_full(tensor: torch.Tensor) -> torch.Tensor:
		279	+ from torch.distributed.tensor import DTensor
		280	+ return tensor.full_tensor().cpu() if isinstance(tensor, DTensor) else tensor
		281	+
277	def dcp_to_hf(	282	def dcp_to_hf(
278	self, key: str, tensor: torch.Tensor	283	self, key: str, tensor: torch.Tensor
279	) -> Dict[str, torch.Tensor]:	284	) -> Dict[str, torch.Tensor]:
	@@ -288,13 +293,13 @@ class Qwen35WeightTransformPipeline(WeightTransformPipeline):
288	# Split merged MTP weights:	293	# Split merged MTP weights:
289	# experts.gate_up_proj [E, H, 2I] -> experts.{e}.{gate,up}_proj.weight [I, H]	294	# experts.gate_up_proj [E, H, 2I] -> experts.{e}.{gate,up}_proj.weight [I, H]
290	# experts.down_proj [E, I, H] -> experts.{e}.down_proj.weight [H, I]	295	# experts.down_proj [E, I, H] -> experts.{e}.down_proj.weight [H, I]
291	- for dcp_key, hf_keys in self.dcp_to_hf_mapping.items():	296	+ hf_keys = self.dcp_to_hf_mapping.get(key)
292	- if key != dcp_key:	297	+ if hf_keys is not None:
293	- continue	298	+ # Gather before unbinding the sharded expert dimension.
294	return split_moe_expert_weights(	299	return split_moe_expert_weights(
295	- tensor=tensor,	300	+ tensor=self._to_full(tensor),
296	hf_keys=hf_keys,	301	hf_keys=hf_keys,
297	- dcp_key=dcp_key,	302	+ dcp_key=key,
298	)	303	)
299		304
300	# Permute fused MoE weights:	305	# Permute fused MoE weights:

Mmindspeed_mm/fsdp/checkpoint/hf_checkpointer.py+17-5

	@@ -24,6 +24,7 @@ from mindspeed_mm.fsdp.checkpoint.hf_utils import (
24	get_model_save_state,	24	get_model_save_state,
25	get_dtype_size,	25	get_dtype_size,
26	find_safetensors_index,	26	find_safetensors_index,
		27	+ merge_lora_weights,
27	save_state_dict,	28	save_state_dict,
28	)	29	)
29	from mindspeed_mm.fsdp.utils.constants import FILE_MODE	30	from mindspeed_mm.fsdp.utils.constants import FILE_MODE
	@@ -46,6 +47,8 @@ class HuggingFaceCheckpointer(CheckpointerBase):
46	model_assets_dir: str = None,	47	model_assets_dir: str = None,
47	model_id=None,	48	model_id=None,
48	enable_lora: bool = False,	49	enable_lora: bool = False,
		50	+ lora_alpha: Optional[int] = None,
		51	+ lora_rank: Optional[int] = None,
49	mtp_num_layers: Optional[int] = None,	52	mtp_num_layers: Optional[int] = None,
50	**kwargs,	53	**kwargs,
51	) -> None:	54	) -> None:
	@@ -59,6 +62,8 @@ class HuggingFaceCheckpointer(CheckpointerBase):
59	model_assets_dir: original huggingface checkpoint directory used to copy config, processor and shard mapping	62	model_assets_dir: original huggingface checkpoint directory used to copy config, processor and shard mapping
60	model_id: model id used to select model-specific weight transform pipeline	63	model_id: model id used to select model-specific weight transform pipeline
61	enable_lora: whether to normalize PEFT base-layer keys before saving	64	enable_lora: whether to normalize PEFT base-layer keys before saving
		65	+ lora_alpha: LoRA scaling numerator used for merged HF export
		66	+ lora_rank: LoRA rank used for merged HF export
62	mtp_num_layers: number of MTP layers enabled by the training model	67	mtp_num_layers: number of MTP layers enabled by the training model
63	return:	68	return:
64	None	69	None
	@@ -76,11 +81,18 @@ class HuggingFaceCheckpointer(CheckpointerBase):
76	fqn_to_filename_mapping = json.load(f)["weight_map"]	81	fqn_to_filename_mapping = json.load(f)["weight_map"]
77	transform_cls = WEIGHT_TRANSFORM_PIPELINES.get(model_id, None)	82	transform_cls = WEIGHT_TRANSFORM_PIPELINES.get(model_id, None)
78	weight_transform = transform_cls(hf_dir=model_assets_dir, mtp_num_layers=mtp_num_layers) if transform_cls else None	83	weight_transform = transform_cls(hf_dir=model_assets_dir, mtp_num_layers=mtp_num_layers) if transform_cls else None
79	- save_state = get_model_save_state(	84	+ save_state = get_model_save_state(model = state["model"], enable_lora=enable_lora)
80	- state["model"],	85	+
81	- save_ckpt_dtype=save_ckpt_dtype,	86	+ # Merge LoRA weights into base weights.
82	- enable_lora=enable_lora,	87	+ if enable_lora:
83	- )	88	+ save_state = merge_lora_weights(save_state, scaling=lora_alpha / lora_rank)
		89	+
		90	+ # Convert floating-point weights to the requested dtype before saving.
		91	+ if save_ckpt_dtype is not None:
		92	+ for key, tensor in save_state.items():
		93	+ if (isinstance(tensor, torch.Tensor) and torch.is_floating_point(tensor)):
		94	+ save_state[key] = tensor.to(save_ckpt_dtype)
		95	+
84	if weight_transform is not None:	96	if weight_transform is not None:
85	new_state = {}	97	new_state = {}
86	for key, tensor in save_state.items():	98	for key, tensor in save_state.items():

Mmindspeed_mm/fsdp/checkpoint/hf_utils.py+44-0

	@@ -11,6 +11,7 @@ from pydantic import FilePath
11	from safetensors import safe_open	11	from safetensors import safe_open
12	from safetensors.torch import save_file	12	from safetensors.torch import save_file
13	import torch	13	import torch
		14	+from torch.distributed.tensor import DTensor, Replicate
14	from tqdm import tqdm	15	from tqdm import tqdm
15	from transformers.utils import SAFE_WEIGHTS_INDEX_NAME, SAFE_WEIGHTS_NAME	16	from transformers.utils import SAFE_WEIGHTS_INDEX_NAME, SAFE_WEIGHTS_NAME
16		17
	@@ -275,6 +276,49 @@ def get_model_save_state(
275	return model_state_cls(model, save_ckpt_dtype).state_dict()	276	return model_state_cls(model, save_ckpt_dtype).state_dict()
276		277
277		278
		279	+@torch.no_grad()
		280	+def merge_lora_weights(
		281	+ state_dict: Dict[str, torch.Tensor],
		282	+ scaling: float,
		283	+) -> Dict[str, torch.Tensor]:
		284	+ """Merge the default LoRA adapter into its base weights.
		285	+
		286	+ For a base linear weight ``W`` with shape ``[out_features, in_features]``,
		287	+ PEFT stores ``lora_A`` with shape ``[rank, in_features]`` and ``lora_B``
		288	+ with shape ``[out_features, rank]``. The adapter update therefore has the
		289	+ same shape as ``W`` and is computed as ``lora_B @ lora_A``.
		290	+
		291	+ FSDP2 can represent ``lora_A`` as a sharded DTensor. The matrix product
		292	+ must use the complete ``[rank, in_features]`` matrix (the rank dimension
		293	+ is needed to contract with ``lora_B``), so ``lora_A`` is redistributed to
		294	+ replicated placement before calculating the update.
		295	+ """
		296	+ suffix_a = ".lora_A.default.weight"
		297	+ suffix_b = ".lora_B.default.weight"
		298	+
		299	+ for lora_a_key in [key for key in state_dict if key.endswith(suffix_a)]:
		300	+ prefix = lora_a_key.removesuffix(suffix_a)
		301	+ lora_b_key = f"{prefix}{suffix_b}"
		302	+ base_key = f"{prefix}.weight"
		303	+
		304	+ base = state_dict[base_key]
		305	+ lora_a = state_dict[lora_a_key]
		306	+ lora_b = state_dict[lora_b_key]
		307	+ if isinstance(lora_a, DTensor):
		308	+ # A sharded lora_A contains only part of the rank dimension on
		309	+ # each rank; aggregate it before the B @ A multiplication.
		310	+ lora_a = lora_a.redistribute(
		311	+ placements=[Replicate() for _ in range(lora_a.device_mesh.ndim)]
		312	+ )
		313	+ delta = torch.matmul(lora_b, lora_a)
		H H htwang9月15日加个注释说明下，这里涉及dtensor和tensor的操作 LKONE 9月15日评论： 9月15日评论：
		314	+
		315	+ state_dict[base_key] = base + scaling * delta
		316	+ del state_dict[lora_a_key]
		317	+ del state_dict[lora_b_key]
		318	+
		319	+ return state_dict
		320	+
		321	+
278	@lru_cache	322	@lru_cache
279	def get_dtype_size(dtype: "torch.dtype") -> int:	323	def get_dtype_size(dtype: "torch.dtype") -> int:
280	_float8_e4m3fn = getattr(torch, "float8_e4m3fn", None)	324	_float8_e4m3fn = getattr(torch, "float8_e4m3fn", None)

Mmindspeed_mm/fsdp/params/lora_args.py+7-0

	@@ -33,6 +33,7 @@ class LoraArguments(BaseArguments):
33		33
34	Attributes:	34	Attributes:
35	enable: Whether to enable LoRA fine-tuning.	35	enable: Whether to enable LoRA fine-tuning.
		36	+ lora_save_only: Whether to export only LoRA weights instead of a full checkpoint.
36	rank: Rank of the low-rank matrices.	37	rank: Rank of the low-rank matrices.
37	alpha: Scaling factor for LoRA weights.	38	alpha: Scaling factor for LoRA weights.
38	target_modules: List of target module names/patterns for LoRA.	39	target_modules: List of target module names/patterns for LoRA.
	@@ -44,6 +45,12 @@ class LoraArguments(BaseArguments):
44	default=False,	45	default=False,
45	metadata={"help": "Enable LoRA fine-tuning."},	46	metadata={"help": "Enable LoRA fine-tuning."},
46	)	47	)
		48	+ lora_save_only: bool = field(
		49	+ default=False,
		50	+ metadata={
		51	+ "help": "When LoRA is enabled, export only adapter safetensors and its config."
		52	+ },
		53	+ )
47	rank: int = field(	54	rank: int = field(
48	default=8,	55	default=8,
49	metadata={"help": "Rank of the low-rank matrices."},	56	metadata={"help": "Rank of the low-rank matrices."},

Mmindspeed_mm/fsdp/train/train_engine.py+16-9

	@@ -445,15 +445,21 @@ class TrainEngine:
445	dataloader_state: Dataloader state aligned with training progress, defaults to None.	445	dataloader_state: Dataloader state aligned with training progress, defaults to None.
446	"""	446	"""
447	args = self.args	447	args = self.args
		448	+ save_ckpt_dtype = (
		449	+ get_dtype(args.training.save_ckpt_dtype)
		450	+ if args.training.save_ckpt_dtype
		451	+ else None
		452	+ )
448		453
449	- # Handle LoRA save modes	454	+ if args.training.lora.enable and args.training.lora.lora_save_only:
450	- if args.training.lora.enable:	455	+ self.lora_weight_manager.save_lora_only(
451	- # Save only LoRA adapter weights	456	+ save_path=args.training.save,
452	- if self.lora_weight_manager is not None:	457	+ iteration=iteration,
453	- self.lora_weight_manager.save_lora_only(	458	+ save_ckpt_dtype=save_ckpt_dtype,
454	- save_path=args.training.save,	459	+ )
455	- iteration=iteration,	460	+ # Synchronize all processes after saving
456	- )	461	+ torch.distributed.barrier()
		462	+ return
457		463
458	if dataloader_state is None:	464	if dataloader_state is None:
459	dataloader_state = self.train_dataloader.state_dict()	465	dataloader_state = self.train_dataloader.state_dict()
	@@ -472,13 +478,14 @@ class TrainEngine:
472	state["optimizer"] = self.optimizer	478	state["optimizer"] = self.optimizer
473	if not args.training.no_save_rng:	479	if not args.training.no_save_rng:
474	state["extra_state"]["torch_rng_state"] = torch.get_rng_state()	480	state["extra_state"]["torch_rng_state"] = torch.get_rng_state()
475	- save_ckpt_dtype = get_dtype(args.training.save_ckpt_dtype) if args.training.save_ckpt_dtype else None
476	self.save_checkpointer.save(	481	self.save_checkpointer.save(
477	args.training.save,	482	args.training.save,
478	state=state,	483	state=state,
479	iteration=iteration,	484	iteration=iteration,
480	save_async=args.training.save_async,	485	save_async=args.training.save_async,
481	enable_lora=args.training.lora.enable,	486	enable_lora=args.training.lora.enable,
		487	+ lora_alpha=args.training.lora.alpha,
		488	+ lora_rank=args.training.lora.rank,
482	save_ckpt_dtype=save_ckpt_dtype,	489	save_ckpt_dtype=save_ckpt_dtype,
483	model_assets_dir=args.model.model_name_or_path,	490	model_assets_dir=args.model.model_name_or_path,
484	model_id=args.model.model_id,	491	model_id=args.model.model_id,

Mmindspeed_mm/fsdp/utils/lora_weight_manager.py+15-4

	@@ -32,6 +32,7 @@ import torch.distributed as dist
32	import torch.nn as nn	32	import torch.nn as nn
33		33
34	from mindspeed.fsdp.utils.log import print_rank	34	from mindspeed.fsdp.utils.log import print_rank
		35	+from mindspeed_mm.fsdp.checkpoint.utils import get_checkpoint_name
35		36
36	try:	37	try:
37	from torch.distributed._tensor import DTensor	38	from torch.distributed._tensor import DTensor
	@@ -108,6 +109,7 @@ class LoraWeightManager:
108	self,	109	self,
109	save_path: str,	110	save_path: str,
110	iteration: Optional[int] = None,	111	iteration: Optional[int] = None,
		112	+ save_ckpt_dtype: Optional[torch.dtype] = None,
111	) -> Tuple[int, int]:	113	) -> Tuple[int, int]:
112	"""Save only LoRA adapter weights.	114	"""Save only LoRA adapter weights.
113		115
	@@ -118,6 +120,7 @@ class LoraWeightManager:
118	Args:	120	Args:
119	save_path: Directory path to save LoRA weights.	121	save_path: Directory path to save LoRA weights.
120	iteration: Optional iteration number for checkpoint naming.	122	iteration: Optional iteration number for checkpoint naming.
		123	+ save_ckpt_dtype: Optional dtype used for the saved adapter weights.
121		124
122	Returns:	125	Returns:
123	Tuple of (num_saved_params, num_lora_params) where:	126	Tuple of (num_saved_params, num_lora_params) where:
	@@ -135,14 +138,22 @@ class LoraWeightManager:
135	"Please install it with: pip install safetensors"	138	"Please install it with: pip install safetensors"
136	) from e	139	) from e
137		140
138	- os.makedirs(save_path, exist_ok=True)	141	+ output_dir = (
		142	+ get_checkpoint_name(save_path, iteration, release=False)
		143	+ if iteration is not None
		144	+ else save_path
		145	+ )
		146	+ os.makedirs(output_dir, exist_ok=True)
139		147
140	lora_state_dict: Dict[str, torch.Tensor] = {}	148	lora_state_dict: Dict[str, torch.Tensor] = {}
141	num_lora_params = 0	149	num_lora_params = 0
142		150
143	for name, param in self.model.named_parameters():	151	for name, param in self.model.named_parameters():
144	if "lora" in name and "base_layer" not in name:	152	if "lora" in name and "base_layer" not in name:
145	- gathered_param = self._gather_dtensor(param.data)	153	+ param_data = param.data
		154	+ if save_ckpt_dtype is not None and torch.is_floating_point(param_data):
		155	+ param_data = param_data.to(save_ckpt_dtype)
		156	+ gathered_param = self._gather_dtensor(param_data)
146	lora_state_dict[name] = gathered_param	157	lora_state_dict[name] = gathered_param
147	num_lora_params += gathered_param.numel()	158	num_lora_params += gathered_param.numel()
148		159
	@@ -151,7 +162,7 @@ class LoraWeightManager:
151	else:	162	else:
152	filename = "lora_adapter.safetensors"	163	filename = "lora_adapter.safetensors"
153		164
154	- save_path_full = os.path.join(save_path, filename)	165	+ save_path_full = os.path.join(output_dir, filename)
155	if self._rank == 0:	166	if self._rank == 0:
156	save_file(lora_state_dict, save_path_full)	167	save_file(lora_state_dict, save_path_full)
157		168
	@@ -159,7 +170,7 @@ class LoraWeightManager:
159		170
160	# Generate a PEFT-compatible adapter_config.json alongside the weights	171	# Generate a PEFT-compatible adapter_config.json alongside the weights
161	# so the output directory can be loaded directly by PEFT/vLLM/SGLang.	172	# so the output directory can be loaded directly by PEFT/vLLM/SGLang.
162	- self.write_adapter_config(save_path)	173	+ self.write_adapter_config(output_dir)
163		174
164	print_rank(	175	print_rank(
165	logger.info,	176	logger.info,

Mtests/ut/test_lora_fsdp2_integration.py+1-1

	@@ -491,7 +491,7 @@ class TestLoraWeightManagerExtended:
491	assert num_saved > 0	491	assert num_saved > 0
492	assert num_elements > 0	492	assert num_elements > 0
493	assert os.path.exists(	493	assert os.path.exists(
494	- os.path.join(save_path, "lora_adapter_iteration_5.safetensors")	494	+ os.path.join(save_path, "iter_0000005", "lora_adapter_iteration_5.safetensors")
495	)	495	)
496		496
497	def test_save_lora_without_iteration(self, tmp_path: pathlib.Path) -> None:	497	def test_save_lora_without_iteration(self, tmp_path: pathlib.Path) -> None: