Test models in FP32, reduce thresholds.

[jlamypoirier]

✨ Description

Fix: #304

New testing features:

• Add fine-grained capacity to comparison configs, so we can vary the thresholds and other options by tensor, step, etc. This will prevent situations where we have to use very high thresholds because of worst-case tensors.
• Add threshold scaling option to model and distributed configs, so we don't have to increase all thresholds because of one config.
• Add option to disable specific distributed configurations for some models that may not support them.
• Show comparison configs (pprint).
• Use a separate directory for each worker when not parallel-safe.
• Delete the entire testing cache on each test run to prevent issues.

Fixes:

• Reduce the vocab size to 384 for model tests. The high vocab size was leading to lots of zero entries in embedding gradients which made comparison difficult ([bug] Model comparison tests are flaky WRT word_embeddings_weight gradients #304)
• Add error for sequence-first SSMs, which would otherwise break the model silently.
• Suppress lots of warnings.

Testing tweaks:

• Run most model tests in fp32 instead of bf16. This allows setting more useful comparison thresholds, and there isn't much specific to bf16 anyway (basically just flash attention).
• Add specific testing configurations for bf16 and fp16, with their own separate comparison thresholds.
• Enable comparison between configs with different hidden state layouts, ex. sequence-first vs not (gradient only).
• Adjust all thresholds so we can use lower values for most comparisons.
• Skip sequence-first tests for SSMs.

🔍 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)