Fix things, prepare for v0.2

[jlamypoirier]

• Fix [bug] Crash with sequence-data-parallel #57
• Fix: [bug] Nans and/or desync for sequence-tensor-parallel. #59 (Kind of. The desync goes away, but no idea about the nans in llama3 that caused that investigation.)

✨ Description

• [Breaking] Drop backward compatibility for configs, except for model configs (so we keep checkpoint compatibility)
• [Breaking] Remove backward compatibility for old state dict checkpoints (we don't have any of these anyway)
• Postpone remaining v0.2 todos to v0.3
• Fix random dataset when the experiment directory is None
• Fix tensor-parallel backup attention
• Fix nan reporting in check_parallel_match
• Fix debug_tensor_parallel for sequence-tensor-parallel
• Fix presents wrongly set for the last micro-batch with sequence-data-parallel, causing a crash ([bug] Crash with sequence-data-parallel #57)
• Add tensor_and_sequence_data distributed dim so we can reconstruct micro-sequences properly (needed for test below)
• Add test for sequence-data-parallel.
• Explicitly shutdown process groups to prevent a warning
• Remove some cross-file test dependencies

🔍 Type of change

Select all that apply:

• 🐛 Bug fix (non-breaking change that addresses a specific issue)
• 🚀 New feature (non-breaking change that adds functionality)
• ⚠️ Breaking change (a change that could affect existing functionality)
• 📈 Performance improvement/optimization (improves speed, memory usage, or efficiency)
• 🛠️ Code refactor (non-functional changes that improve code readability, structure, etc.)
• 📦 Dependency bump (updates dependencies, including Dockerfile or package changes)
• 📝 Documentation change (updates documentation, including new content or typo fixes)
• 🔧 Infrastructure/Build change (affects build process, CI/CD, or dependencies)

[tscholak]

@charlesGE can you confirm that we don't need this anymore please?

[tscholak]

+1

[tscholak]

LGTM, but I'd like to defer to @charlesGE who should try to work with ghcr.io/servicenow/fast-llm:sha-85dfeea and see if our experiments can be continued with it.