shank commited on
Commit ·
cf25957
1
Parent(s): 2c50d8a
Revert speed optimizations to prioritize real-world model quality
Browse files- training/train_grpo.py +5 -5
training/train_grpo.py
CHANGED
|
@@ -34,7 +34,7 @@ parser.add_argument("--test", action="store_true", help="Run 10 steps for testin
|
|
| 34 |
parser.add_argument("--test-local", action="store_true", dest="test_local",
|
| 35 |
help="Sanity-check reward function locally without any model or GPU")
|
| 36 |
parser.add_argument("--resume", type=str, default=None, help="Path to checkpoint")
|
| 37 |
-
parser.add_argument("--max_steps", type=int, default=
|
| 38 |
args = parser.parse_args()
|
| 39 |
|
| 40 |
|
|
@@ -165,9 +165,9 @@ def load_bugs(tier: int) -> list[dict]:
|
|
| 165 |
|
| 166 |
def get_bugs_for_step(step: int) -> list[dict]:
|
| 167 |
tier1 = load_bugs(1)
|
| 168 |
-
if step <
|
| 169 |
return tier1
|
| 170 |
-
elif step <
|
| 171 |
return tier1 + load_bugs(2)
|
| 172 |
return tier1 + load_bugs(2) + load_bugs(3)
|
| 173 |
|
|
@@ -328,7 +328,7 @@ elif _gpu_vram_gb >= 20: # A10G 24GB / V100 32GB
|
|
| 328 |
_lora_r = 8
|
| 329 |
else: # T4 15GB / anything smaller
|
| 330 |
_batch = 2
|
| 331 |
-
_grad_accum =
|
| 332 |
_num_gen = 2 # 2 % 2 == 0
|
| 333 |
_max_comp = 160
|
| 334 |
_lora_r = 8
|
|
@@ -509,7 +509,7 @@ trainer = GRPOTrainer(
|
|
| 509 |
class CurriculumCallback(TrainerCallback):
|
| 510 |
def on_step_end(self, args, state, control, **kwargs):
|
| 511 |
step = state.global_step
|
| 512 |
-
if step in [
|
| 513 |
trainer.train_dataset = make_dataset(step)
|
| 514 |
print(f"\nCurriculum advanced at step {step}!")
|
| 515 |
if WANDB_API_KEY:
|
|
|
|
| 34 |
parser.add_argument("--test-local", action="store_true", dest="test_local",
|
| 35 |
help="Sanity-check reward function locally without any model or GPU")
|
| 36 |
parser.add_argument("--resume", type=str, default=None, help="Path to checkpoint")
|
| 37 |
+
parser.add_argument("--max_steps", type=int, default=500)
|
| 38 |
args = parser.parse_args()
|
| 39 |
|
| 40 |
|
|
|
|
| 165 |
|
| 166 |
def get_bugs_for_step(step: int) -> list[dict]:
|
| 167 |
tier1 = load_bugs(1)
|
| 168 |
+
if step < 150:
|
| 169 |
return tier1
|
| 170 |
+
elif step < 600:
|
| 171 |
return tier1 + load_bugs(2)
|
| 172 |
return tier1 + load_bugs(2) + load_bugs(3)
|
| 173 |
|
|
|
|
| 328 |
_lora_r = 8
|
| 329 |
else: # T4 15GB / anything smaller
|
| 330 |
_batch = 2
|
| 331 |
+
_grad_accum = 4
|
| 332 |
_num_gen = 2 # 2 % 2 == 0
|
| 333 |
_max_comp = 160
|
| 334 |
_lora_r = 8
|
|
|
|
| 509 |
class CurriculumCallback(TrainerCallback):
|
| 510 |
def on_step_end(self, args, state, control, **kwargs):
|
| 511 |
step = state.global_step
|
| 512 |
+
if step in [150, 350]:
|
| 513 |
trainer.train_dataset = make_dataset(step)
|
| 514 |
print(f"\nCurriculum advanced at step {step}!")
|
| 515 |
if WANDB_API_KEY:
|