Error due to missing is_zero arg when saving LR scheduler

[Lauler]

This recent commit 51bd072 and pull request #230 which changes how LR schedulers are saved forgot to add the is_zero arg and set it to config.optimizer.zero_stage in src/nanotron/serialize/main.py:

nanotron/src/nanotron/serialize/main.py

Lines 106 to 110 in cfcdeae

	save_lr_scheduler(
	lr_scheduler=lr_scheduler,
	parallel_context=parallel_context,
	root_folder=root_folder,
	)

This arg is expected to be passed:

nanotron/src/nanotron/serialize/optimizer.py

Lines 109 to 114 in cfcdeae

	def save_lr_scheduler(
	lr_scheduler,
	is_zero,
	parallel_context: ParallelContext,
	root_folder: Path,
	):

This causes training to crash when LR scheduler is saved.

@TJ-Solergibert @NouamaneTazi

[sankexin]

nanotron/src/nanotron/serialize/main.py

save_lr_scheduler(
    lr_scheduler=lr_scheduler,
    is_zero=True,
    parallel_context=parallel_context,
    root_folder=root_folder,
)

[Lauler]

I changed mine to

save_lr_scheduler( 
     lr_scheduler=lr_scheduler,
     is_zero=config.optimizer.zero_stage,
     parallel_context=parallel_context, 
     root_folder=root_folder, 
 )