Skip to content

Commit d0dcc32

Browse files
daniellepintzfacebook-github-bot
authored andcommitted
Add qwen3_30b_a3b.yaml config for test_policy_update.py (#682)
Summary: - Add a config to test weight sync for Qwen3 30B MoE model - Fix docstring in test_policy_update.py Reviewed By: allenwang28 Differential Revision: D89833611
1 parent 462de97 commit d0dcc32

File tree

2 files changed

+83
-3
lines changed

2 files changed

+83
-3
lines changed
Lines changed: 80 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,80 @@
1+
# pytest -s tests/integration_tests/test_policy_update.py::TestWeightSync::test_sanity_check --config tests/integration_tests/fixtures/qwen3_30b_a3b.yaml
2+
3+
# trainer tp = 1, generator tp = 4
4+
5+
# Global configuration
6+
group_size: 8
7+
batch_size: 2
8+
max_req_tokens: 256
9+
max_res_tokens: 256
10+
model: "Qwen/Qwen3-30B-A3B"
11+
off_by_n: 1 # Off by one by default
12+
compile: false # Enable torch.compile for trainer, and CUDA graphs for vLLM
13+
14+
15+
# Generator configuration
16+
generator:
17+
engine_args:
18+
model: ${model}
19+
tensor_parallel_size: 4
20+
pipeline_parallel_size: 1
21+
enforce_eager: ${not:${compile}}
22+
sampling_params:
23+
n: ${group_size}
24+
max_tokens: ${max_res_tokens}
25+
temperature: 1.0
26+
top_p: 1.0
27+
28+
# Trainer configuration
29+
trainer:
30+
model:
31+
name: qwen3
32+
flavor: 30B-A3B
33+
hf_assets_path: hf://${model}
34+
optimizer:
35+
name: AdamW
36+
lr: 1e-5
37+
eps: 1e-8
38+
lr_scheduler:
39+
warmup_steps: 1
40+
training:
41+
local_batch_size: ${batch_size}
42+
seq_len: ${sum:${max_req_tokens},${max_res_tokens}} # seq_len >= max_req_tokens + max_res_tokens
43+
max_norm: 1.0
44+
steps: 1000000
45+
dtype: bfloat16
46+
gc_freq: 1
47+
compile:
48+
enable: ${compile}
49+
parallelism:
50+
data_parallel_replicate_degree: 1
51+
data_parallel_shard_degree: -1
52+
tensor_parallel_degree: 1
53+
pipeline_parallel_degree: 1
54+
context_parallel_degree: 1
55+
expert_parallel_degree: 1
56+
expert_tensor_parallel_degree: 1
57+
disable_loss_parallel: true
58+
checkpoint:
59+
enable: true
60+
initial_load_path: hf://${model}
61+
initial_load_in_hf: true
62+
last_save_in_hf: true
63+
interval: 500
64+
async_mode: "disabled"
65+
activation_checkpoint:
66+
mode: selective
67+
selective_ac_option: op
68+
69+
# All resource allocations
70+
services:
71+
generator:
72+
procs: ${generator.engine_args.tensor_parallel_size}
73+
num_replicas: 1
74+
with_gpus: true
75+
76+
actors:
77+
trainer:
78+
procs: 4
79+
num_replicas: 1
80+
with_gpus: true

tests/integration_tests/test_policy_update.py

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -210,9 +210,9 @@ async def test_sanity_check(self, _setup_and_teardown):
210210
211211
The check performs the following steps:
212212
- Initialize trainer and push weights v0 (original huggingface ckpt)
213-
- Step the trainer, setting all weights to zero and push weights v1
214-
- Load weights v0 and check the generator has all zero weights
215-
- Load weights v1 and check the generator has all the weights back
213+
- Set all the weights of the model on the trainer to zero and push weights as v1
214+
- Load weights v1 and check the generator has all zero weights
215+
- Load weights v0 and check the generator has all the original weights back
216216
217217
"""
218218

0 commit comments

Comments
 (0)