Merge pull request #439 from Modalities/nccl_compute_overlap_investigation

chore: removed comments

chore: added documentatoin of linear warmup cosine annealing learning rate scheduler

feat: added LinearWarmupCosineAnnealingLRScheduler

refactor: reintroducing layers_per_fsdp_unit in model_factory for better FSDP sharding control

chore: setting cycle_momentum to False in OneCycleLRSchedulerConfig

feat: added garbage collection

Merge pull request #440 from Modalities/layer_norm

chore: Move parameter norm scripts to separate folder

feat: Group plots by layer

feat: Compute norms for parameters instead of layers

refactor: Separate logic for layer norm computation and plotting

feat: Support multiple checkpoints

feat: Compute layer norms for a model checkpoint

chore: temporarily reverted combination of fsdp units