fix(train): consider schedule offset

borisdayma · borisdayma · commit bc4734ff3f83 · 2022-02-23T04:37:57.000Z
diff --git a/tools/train/train.py b/tools/train/train.py
@@ -688,7 +688,8 @@ def create_learning_rate_fn() -> Callable[[int], jnp.array]:
                 staircase=training_args.lr_staircase,
             )
         schedule_fn = optax.join_schedules(
-            schedules=[warmup_fn, decay_fn], boundaries=[training_args.warmup_steps]
+            schedules=[warmup_fn, decay_fn],
+            boundaries=[model_metadata.get("step", 0) + training_args.warmup_steps],
         )
         return schedule_fn
 

Original file line number	Diff line number	Diff line change
`@@ -688,7 +688,8 @@ def create_learning_rate_fn() -> Callable[[int], jnp.array]:`
`688`	`688`	`staircase=training_args.lr_staircase,`
`689`	`689`	`)`
`690`	`690`	`schedule_fn = optax.join_schedules(`
`691`		`- schedules=[warmup_fn, decay_fn], boundaries=[training_args.warmup_steps]`
	`691`	`+ schedules=[warmup_fn, decay_fn],`
	`692`	`+ boundaries=[model_metadata.get("step", 0) + training_args.warmup_steps],`
`692`	`693`	`)`
`693`	`694`	`return schedule_fn`
`694`	`695`