Merge pull request #384 from Modalities/fix/warmstart_fsdp2_weight_tying

Migration to latest versions of torch & flash-attn to solve warmstart/fsdp2/weight tying problem

Author: le1nux

.github/workflows/build_and_deploy_documentation.yml

@@ -25,7 +25,7 @@ jobs:
       run: |
         sudo apt-get update
         sudo apt-get install git -y
-        python -m pip install torch==2.6.0
+        python -m pip install torch==2.7.1
         python -m pip install --upgrade pip setuptools wheel
         export FLASH_ATTENTION_SKIP_CUDA_BUILD=TRUE
         python -m pip install -e .

.github/workflows/tests_full.yml

@@ -23,11 +23,11 @@ jobs:
         sudo apt-get update
         sudo apt-get install curl -y                  # required by coveralls
         sudo apt-get install git -y
-        python -m pip install torch==2.6.0
+        python -m pip install torch==2.7.1
         python -m pip install --upgrade pip setuptools wheel
         export FLASH_ATTENTION_SKIP_CUDA_BUILD=TRUE
         python -m pip install ninja                   # Lowers compilation time of flash attention significantly 
-        python -m pip install flash-attn==2.7.4.post1 --no-build-isolation
+        python -m pip install flash-attn==2.8.0.post2 --no-build-isolation
         python -m pip install -e .[tests]
     - name: Run tests
       run: |

README.md

@@ -44,9 +44,9 @@ conda create -n modalities python=3.11
 conda activate modalities
 
 # install PyTorch, Ninja and Flash Attention (mandatory)
-pip install torch==2.6.0
+pip install torch==2.7.1
 pip install ninja     # Lowers compilation time of flash attention significantly 
-pip install flash-attn==2.7.4.post1 --no-build-isolation
+pip install flash-attn==2.8.0.post2 --no-build-isolation
 ```
 
 ### Option 1: Installation from source
@@ -82,9 +82,9 @@ pip install modalities
 curl -LsSf https://astral.sh/uv/install.sh | sh
 uv venv --seed --python 3.11 --prompt modalities
 source .venv/bin/activate
-uv pip install torch
+uv pip install torch==2.7.1
 uv pip install ninja
-uv pip install --no-build-isolation flash-attn==2.7.4.post1
+uv pip install --no-build-isolation flash-attn==2.8.0.post2
 # for developer: use [tests,linting] and install pre-commit hooks
 uv pip install -e .[tests,linting]
 pre-commit install --install-hooks

config_files/training/config_lorem_ipsum_long_fsdp1_warmstart.yaml

@@ -68,7 +68,7 @@ settings:
       variant_key: last_step_from_checkpoint_path
       config:
         checkpoint_path: ${settings.warmstart_checkpoint_paths.model_checkpoint_path}
-  warmstart_checkpoint_paths: ${warmstart_env:checkpoint_paths}
+  warmstart_checkpoint_paths: ${warmstart_env:checkpoint_paths}  # use modalities warmstart [..] --last_checkpoint_info_file_path [..]
 
 collate_fn:
   component_key: collate_fn

config_files/training/config_lorem_ipsum_long_fsdp2_warmstart.yaml

@@ -74,8 +74,7 @@ settings:
       variant_key: last_step_from_checkpoint_path
       config:
         checkpoint_path: ${settings.warmstart_checkpoint_paths.checkpoint_folder_path}
-  warmstart_checkpoint_paths: # ${warmstart_env:checkpoint_paths}
-    checkpoint_folder_path: /raid/s3/opengptx/max_lue/repositories/modalities/data/checkpoints/2025-03-14__15-25-59_970fedec/eid_2025-03-14__15-25-59_970fedec-seen_steps_96-seen_tokens_196608-target_steps_162-target_tokens_331776
+  warmstart_checkpoint_paths: ${warmstart_env:checkpoint_paths}  # use modalities warmstart [..] --last_checkpoint_info_file_path [..]
 
 collate_fn:
   component_key: collate_fn

src/modalities/logging_broker/subscriber_impl/results_subscriber.py

@@ -73,7 +73,12 @@ def __init__(
         with open(config_file_path, "r", encoding="utf-8") as file:
             config = yaml.safe_load(file)
         self.run = wandb.init(
-            project=project, name=experiment_id, mode=mode.value.lower(), dir=logging_directory, config=config
+            project=project,
+            name=experiment_id,
+            mode=mode.value.lower(),
+            dir=logging_directory,
+            config=config,
+            settings=wandb.Settings(init_timeout=120),
         )
 
         self.run.log_artifact(config_file_path, name=f"config_{wandb.run.id}", type="config")

tutorials/warmstart/configs/warmstart_config.yaml

@@ -53,7 +53,7 @@ settings:
       variant_key: last_step_from_checkpoint_path
       config:
         checkpoint_path: ${settings.warmstart_checkpoint_paths.checkpoint_folder_path}
-  warmstart_checkpoint_paths: ${warmstart_env:checkpoint_paths}
+  warmstart_checkpoint_paths: ${warmstart_env:checkpoint_paths}  # use modalities warmstart [..] --last_checkpoint_info_file_path [..]
 
 collate_fn:
   component_key: collate_fn