chore: Merge branch 'fsdp2_min_integration' into fix_failing_test_for_fsdp2_integration

Author: le1nux

README.md

@@ -94,7 +94,7 @@ Example:
 CUDA_VISIBLE_DEVICES=0,1,2,3 torchrun --rdzv-endpoint localhost:29515 \
                                         --nnodes 1 \
                                         --nproc_per_node 4 \
-                                        $(which modalities) run --config_file_path configs/pretraining_config.yaml
+                                        $(which modalities) run --config_file_path config_files/training/config_lorem_ipsum_long_fsdp2.yaml
 ```
 
 Explanation:
@@ -111,7 +111,7 @@ Explanation:
 
 * `$(which modalities) run`: This part dynamically finds the path to the Modalities executable and runs it. The run command triggers the main process to start the training.
 
-* `--config_file_path configs/pretraining_config.yaml`: The --config_file_path argument provides the path to the configuration file for the training job. In the example above, it is given by `configs/pretraining_config.yaml`. A configuraton file contains an exhaustive parameterization for all the training components (e.g., dataset, model, optimizer, etc.), making training fully reproducible. An example configuration file can be found [here](tutorials/getting_started/example_config.yaml), and a complete list of components available in Modalities is provided [here](docs/components/components.md).
+* `--config_file_path config_files/training//config_lorem_ipsum_long_fsdp2.yaml`: The --config_file_path argument provides the path to the configuration file for the training job. In the example above, it is given by `config_files/training/config_lorem_ipsum_long_fsdp2.yaml`. A configuraton file contains an exhaustive parameterization for all the training components (e.g., dataset, model, optimizer, etc.), making training fully reproducible. An example configuration file can be found [here](tutorials/getting_started/example_config.yaml), and a complete list of components available in Modalities is provided [here](docs/components/components.md).
 
 If you are a VSCode user, you may want to add this to your `launch.json`:
 ```json

config_files/training/config_example_coca.yaml

@@ -182,9 +182,23 @@ loss_fn:
     target_key: ${settings.referencing_keys.target_key}
     prediction_key: ${settings.referencing_keys.prediction_key}
 
+app_state:
+  component_key: app_state
+  variant_key: raw
+  config:
+    model: 
+      instance_key: wrapped_model
+      pass_type: BY_REFERENCE
+    optimizer:
+      instance_key: optimizer
+      pass_type: BY_REFERENCE
+    lr_scheduler:
+      instance_key: lr_scheduler
+      pass_type: BY_REFERENCE
+
 wrapped_model:
   component_key: model
-  variant_key: fsdp_wrapped
+  variant_key: fsdp1_wrapped
   config:
     model:
       instance_key: model
@@ -256,7 +270,7 @@ model_raw:
     bias_attn_pool: False
     epsilon_attn_pool: 1e-5
 
-scheduler:
+lr_scheduler:
   component_key: scheduler
   variant_key: onecycle_lr
   config:
@@ -286,7 +300,7 @@ optimizer:
 
 gradient_clipper:
   component_key: gradient_clipper
-  variant_key: fsdp_logging_only
+  variant_key: fsdp1_logging_only
   config:
     wrapped_model:
       instance_key: wrapped_model