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

[flxst]

What does this PR do?

This PR addresses #381. The problem can be traced back to a bug in torch 2.6 related to the fact that we flatten the optimizer state dict here and here.

A solution is to simply migrate to torch 2.7. This requires to also migrate flash-attn to version 2.8.

This PR includes both migrations, along with minor adjustments of the warmstart config files.

Unit tests pass with github actions.

General Changes

• None

Breaking Changes

• None

Checklist before submitting final PR

• My PR is minimal and addresses one issue in isolation
• I have merged the latest version of the target branch into this feature branch
• I have reviewed my own code w.r.t. correct implementation, missing type hints, proper documentation, etc.
• I have run a sample config for model training
• I have checked that all tests run through (python tests/tests.py)
• I have updated the internal changelog (CHANGELOG_DEV.md)

[flxst]

Everything seems to work if conda is used. However, with uv or python -m venv, the installation of flash-attn==2.8.0.post2 fails, as reported here.

[flxst]

Follow-up problem: Dao-AILab/flash-attention#1708

• torch==2.6.0 & flash-attn==2.7.4.post1: works
• torch==2.7.1 & flash-attn==2.8.0.post2: sometimes fails (depending on platform)