[bug] `test_checkpoint` test not passing when any lr scale is set to 0

[oleksost]

🐞 Describe the Bug

Whenever the lr scale of any component is set to 0, e.g. model.base_model.transformer.mlp_lr_scale=0, test_checkpoint is failing with:

FAILED tests/test_checkpoint.py::test_load_pretrained_distributed_checkpoint - AssertionError: torch.Size([0]) != torch.Size([786432])

I wonder how critical is this for loading/saving checkpoints that were trained with lr scaling?
Maybe related to #256.

🔄 Steps to Reproduce

Steps to reproduce the behavior:
add e.g. model.base_model.transformer.mlp_lr_scale=0 here and run test_checkpoint.

Same is the case when lr is set to zero using per layer lr scale from #243 and #258 (yet in this case more than just one test in test_checkpoint fail.
Importantly, if the line self.requires_grad = requires_grad and any(lr_scale_ != 0 for lr_scale_ in self.lr_scale) here is replaced with simple self.requires_grad = requires_grad the test passes.

🎯 Expected Behavior

Test passes.

📜 Environment Information

📝 Additional Context

[jlamypoirier]

@RaymondLi0 had a similar bug, I think it had something to do with debug prints?

[RaymondLi0]

There seems to be some issues with frozen weights:
1 - the debug print here

Fast-LLM/fast_llm/engine/multi_stage/stage.py

Line 159 in 3ac976b

if self._config.debug_all_param_gradients:

tries to log gradients for all params and fails for frozen weights. Would make sense to only log params that require_grad?

2 - when loading from distributed format, a validation fails for frozen weights:

Fast-LLM/fast_llm/engine/multi_stage/fsdp.py

Line 255 in 3ac976b

self._weight_shard_meta.validate(shard)

3 - when loading from distributed format, involving frozen weights either in the loaded-config, or current-config there are issues here

Fast-LLM/fast_llm/engine/multi_stage/fsdp.py

Line 455 in 3ac976b

shard[begin:end][overlap_mask] = loaded_shards[shard_name][overlap_index_map_masked]

I'm not sure what to do here. We could skip this when the current shard is empty. When the loaded shard is empty, should we set the current-shard to zero?

[RaymondLi0]

Proposed some fixes in #270 based on my understanding of the frozen-weights implementation. Please take a look @jlamypoirier