Skip to content

Commit 097678a

Browse files
RayenTianRayenTian
committed
Moved LoRA configuration
Signed-off-by: ruit <ruit@nvidia.com>
1 parent 1321ac4 commit 097678a

6 files changed

Lines changed: 60 additions & 55 deletions

File tree

‎docs/guides/sft.md‎

Lines changed: 19 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -167,22 +167,27 @@ Upon completion of the training process, you can refer to our [evaluation guide]
167167

168168
NeMo RL supports LoRA (Low-Rank Adaptation) for parameter-efficient fine-tuning. LoRA reduces trainable parameters by using low-rank matrices for weight updates while keeping the base model frozen.
169169

170+
Notes:
171+
- LoRA is supported with DTensor v2 and Megatron backends. DTensor v1 does not support LoRA (ensure `policy.dtensor_cfg._v2=true` when using DTensor).
172+
- Triton kernels are only used in the DTensor v2 path. For TP > 1, Automodel currently does not support Triton kernels (see note below).
173+
170174
### Configuration Parameters
171175

172-
The LoRA configuration is specified under the `policy.lora_cfg` section:
176+
The LoRA configuration is specified under the `policy.dtensor_cfg.lora_cfg` section:
173177

174178
policy:
175-
lora_cfg:
176-
enabled: False # Set to True to enable LoRA fine-tuning
177-
target_modules: [] # List of module names to apply LoRA
178-
exclude_modules: [] # List of module names to exclude from LoRA
179-
match_all_linear: true # Apply LoRA to all linear layers
180-
dim: 8 # LoRA rank (r): controls adaptation capacity
181-
alpha: 32 # LoRA scaling factor (effective lr = alpha/dim)
182-
dropout: 0.0 # Dropout probability for LoRA layers
183-
dropout_position: "post" # Dropout position: "pre" or "post"
184-
lora_A_init: "xavier" # Initialization method: "xavier" or "uniform"
185-
use_triton: true # Use Triton-optimized kernels
179+
dtensor_cfg:
180+
lora_cfg:
181+
enabled: False # Set to True to enable LoRA fine-tuning
182+
target_modules: [] # List of module names to apply LoRA
183+
exclude_modules: [] # List of module names to exclude from LoRA
184+
match_all_linear: true # Apply LoRA to all linear layers
185+
dim: 8 # LoRA rank (r): controls adaptation capacity
186+
alpha: 32 # LoRA scaling factor (effective lr = alpha/dim)
187+
dropout: 0.0 # Dropout probability for LoRA layers
188+
dropout_position: "post" # Dropout position: "pre" or "post"
189+
lora_A_init: "xavier" # Initialization method: "xavier" or "uniform"
190+
use_triton: true # Use Triton-optimized kernels (DTensor v2 path)
186191

187192
### Parameter Details
188193
- **`enabled`** (bool): Whether to enable LoRA training
@@ -194,12 +199,12 @@ policy:
194199
- **`dropout`** (float): Dropout probability for regularization
195200
- **`dropout_position`** (str): Apply dropout before ("pre") or after ("post") LoRA
196201
- **`lora_A_init`** (str): Initialization method for LoRA A matrix
197-
- **`use_triton`** (bool): Use Triton-optimized kernels for better performance. Used for dtensor v2 only. **Note**: [Automodel not support triton for TP > 1](https://github.com/NVIDIA-NeMo/Automodel/blob/b2db55eee98dfe81a8bfe5e23ac4e57afd8ab261/nemo_automodel/recipes/llm/train_ft.py#L199). Set to `false` when `tensor_parallel_size > 1` to avoid compatibility issues
202+
- **`use_triton`** (bool): Use Triton-optimized kernels for better performance. Used for DTensor v2 only. **Note**: [Automodel does not support Triton for TP > 1](https://github.com/NVIDIA-NeMo/Automodel/blob/b2db55eee98dfe81a8bfe5e23ac4e57afd8ab261/nemo_automodel/recipes/llm/train_ft.py#L199). Set to `false` when `tensor_parallel_size > 1` to avoid compatibility issues
198203

199204
### Example Usage
200205

201206
```bash
202-
uv run examples/run_sft.py policy.lora_cfg.enabled=True
207+
uv run examples/run_sft.py policy.dtensor_cfg.lora_cfg.enabled=true
203208
```
204209

205210
For more details on LoRA, see [LoRA: Low-Rank Adaptation of Large Language Models](https://arxiv.org/abs/2106.09685).

‎examples/configs/recipes/llm/sft-llama3.1-8b-1n8g-fsdp2tp1-lora.yaml.disabled‎

Lines changed: 5 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -12,12 +12,13 @@ policy:
1212
tokenizer:
1313
name: meta-llama/Llama-3.1-8B-Instruct
1414
chat_template: default
15+
dtensor_cfg:
16+
lora_cfg:
17+
enabled: true
18+
dim: 128
19+
alpha: 128
1520
train_global_batch_size: 128
1621
max_total_sequence_length: 4096
17-
lora_cfg:
18-
enabled: true
19-
dim: 128
20-
alpha: 128
2122
make_sequence_length_divisible_by: 2
2223
optimizer:
2324
kwargs:

‎examples/configs/sft.yaml‎

Lines changed: 12 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -46,18 +46,18 @@ policy:
4646
context_parallel_size: 1
4747
custom_parallel_plan: null
4848

49-
# LoRA (Low-Rank Adaptation) Configuration
50-
lora_cfg:
51-
enabled: False # Set to True to enable LoRA fine-tuning
52-
target_modules: [] # List of module names to apply LoRA (empty list with match_all_linear=true applies to all linear layers)
53-
exclude_modules: [] # List of module names to exclude from LoRA
54-
match_all_linear: true # If True, applies LoRA to all linear layers (overrides target_modules)
55-
dim: 8 # LoRA rank (r): lower rank = fewer parameters but less capacity. Typical values: 4, 8, 16, 32, 64
56-
alpha: 32 # LoRA scaling factor: effective learning rate multiplier = alpha/dim. Typical values: 16, 32, 64
57-
dropout: 0.0 # Dropout probability applied to LoRA layers (0.0 = no dropout)
58-
dropout_position: "post" # Where to apply dropout: "pre" (before LoRA) or "post" (after LoRA)
59-
lora_A_init: "xavier" # Initialization method for LoRA A matrix: "xavier" or "uniform"
60-
use_triton: true # Use Triton-optimized kernels for LoRA (faster but requires flash-attn). Disable when tensor_parallel_size > 1
49+
# LoRA (Low-Rank Adaptation) Configuration
50+
lora_cfg:
51+
enabled: False # Set to True to enable LoRA fine-tuning
52+
target_modules: [] # List of module names to apply LoRA (empty list with match_all_linear=true applies to all linear layers)
53+
exclude_modules: [] # List of module names to exclude from LoRA
54+
match_all_linear: true # If True, applies LoRA to all linear layers (overrides target_modules)
55+
dim: 8 # LoRA rank (r): lower rank = fewer parameters but less capacity. Typical values: 4, 8, 16, 32, 64
56+
alpha: 32 # LoRA scaling factor: effective learning rate multiplier = alpha/dim. Typical values: 16, 32, 64
57+
dropout: 0.0 # Dropout probability applied to LoRA layers (0.0 = no dropout)
58+
dropout_position: "post" # Where to apply dropout: "pre" (before LoRA) or "post" (after LoRA)
59+
lora_A_init: "xavier" # Initialization method for LoRA A matrix: "xavier" or "uniform"
60+
use_triton: true # Use Triton-optimized kernels for LoRA (faster but requires flash-attn). Disable when tensor_parallel_size > 1
6161

6262
dynamic_batching:
6363
enabled: false

‎nemo_rl/models/policy/__init__.py‎

Lines changed: 18 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,23 @@
1717
from nemo_rl.models.generation.interfaces import GenerationConfig
1818

1919

20+
class LoRAConfigDisabled(TypedDict):
21+
enabled: Literal[False]
22+
23+
24+
class LoRAConfig(TypedDict):
25+
enabled: Literal[True]
26+
target_modules: list[str]
27+
exclude_modules: list[str]
28+
match_all_linear: NotRequired[bool]
29+
dim: int
30+
alpha: int
31+
dropout: float
32+
dropout_position: Literal["pre", "post"]
33+
lora_A_init: str
34+
use_triton: NotRequired[bool]
35+
36+
2037
class DTensorConfigDisabled(TypedDict):
2138
enabled: Literal[False]
2239

@@ -32,6 +49,7 @@ class DTensorConfig(TypedDict):
3249
context_parallel_size: int
3350
custom_parallel_plan: str | None
3451
clear_cache_every_n_steps: NotRequired[int | None]
52+
lora_cfg: NotRequired[LoRAConfig | LoRAConfigDisabled]
3553

3654

3755
class SequencePackingConfigDisabled(TypedDict):
@@ -131,23 +149,6 @@ class MegatronConfig(TypedDict):
131149
distributed_data_parallel_config: MegatronDDPConfig
132150

133151

134-
class LoRAConfigDisabled(TypedDict):
135-
enabled: Literal[False]
136-
137-
138-
class LoRAConfig(TypedDict):
139-
enabled: Literal[True]
140-
target_modules: list[str]
141-
exclude_modules: list[str]
142-
match_all_linear: NotRequired[bool]
143-
dim: int
144-
alpha: int
145-
dropout: float
146-
dropout_position: Literal["pre", "post"]
147-
lora_A_init: str
148-
use_triton: NotRequired[bool]
149-
150-
151152
class TokenizerConfig(TypedDict):
152153
name: str
153154
chat_template: NotRequired[str]
@@ -206,7 +207,6 @@ class PolicyConfig(TypedDict):
206207
reward_model_cfg: NotRequired[RewardModelConfig]
207208
dtensor_cfg: DTensorConfig | DTensorConfigDisabled
208209
megatron_cfg: NotRequired[MegatronConfig | MegatronConfigDisabled]
209-
lora_cfg: NotRequired[LoRAConfig | LoRAConfigDisabled]
210210
hf_config_overrides: NotRequired[dict[str, Any]]
211211
dynamic_batching: DynamicBatchingConfig | DynamicBatchingConfigDisabled
212212
sequence_packing: NotRequired[SequencePackingConfig | SequencePackingConfigDisabled]

‎nemo_rl/models/policy/lm_policy.py‎

Lines changed: 5 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -112,12 +112,11 @@ def __init__(
112112
if use_v2:
113113
worker_builder_cls = "nemo_rl.models.policy.workers.dtensor_policy_worker_v2.DTensorPolicyWorkerV2"
114114
else:
115-
assert config.get("lora_cfg", {}).get("enabled", False) is False, (
116-
"LoRA is not supported for DTensorPolicyWorker V1"
117-
)
118-
worker_builder_cls = (
119-
"nemo_rl.models.policy.workers.dtensor_policy_worker.DTensorPolicyWorker"
120-
)
115+
assert (
116+
config["dtensor_cfg"].get("lora_cfg", {}).get("enabled", False)
117+
is False
118+
), "LoRA is not supported for DTensorPolicyWorker V1"
119+
worker_builder_cls = "nemo_rl.models.policy.workers.dtensor_policy_worker.DTensorPolicyWorker"
121120

122121
tp_size = config["dtensor_cfg"]["tensor_parallel_size"]
123122
cp_size = config["dtensor_cfg"]["context_parallel_size"]

‎nemo_rl/models/policy/workers/dtensor_policy_worker_v2.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -239,7 +239,7 @@ def __init__(
239239
model_state_dict_keys = None
240240

241241
# lora config
242-
lora_cfg = self.cfg.get("lora_cfg", None)
242+
lora_cfg = self.cfg["dtensor_cfg"].get("lora_cfg", None)
243243
self.peft_config = None
244244
self.lora_enabled = lora_cfg is not None and lora_cfg["enabled"]
245245
# patch the init_lora_weights method to use the xavier initialization

0 commit comments

Comments
 (0)