[Bug] DionConfig incorrectly scales LM head learning rate for AdamW + Proposal for Dion3 integration

### Summary
In `olmo_core/optim/dion.py` (line 98), the default LM head override divides the learning rate by `math.sqrt(model_dim)` when setting up the AdamW group:

```python
lm_head_override = OptimGroupOverride(
params=params["lm_head"],
opts=dict(algorithm="adamw", lr=self.lr / math.sqrt(model_dim)),
)
```

### Problem & Upstream Discrepancy

According to the reference implementation in `microsoft/dion` (`train.py`):

```python
def lm_head_lr_scale(scalar_opt: str, model_dim: int) -> float:
# See Dion paper App. D.2 and Fig. 14 for the Lion-specific 1 / sqrt(d_in) scale.
return 1 / math.sqrt(model_dim) if scalar_opt == "lion" else 1.0

```

The $1/\sqrt{d_{\text{in}}}$ scaling factor is explicitly Lion-specific because Lion uses `sign()` updates and cannot adjust for the growing gradient variance across wider embeddings. (See [Dion paper](https://arxiv.org/html/2504.05295v2#A4).)

Because OLMo uses **AdamW** for `lm_head`, the second-moment accumulator ($\sqrt{v}$) already normalizes the update magnitude. Applying `self.lr / math.sqrt(model_dim)` double-penalizes the layer, collapsing the effective learning rate (e.g., from `3e-4` down to `~8.4e-6` on a 1280-dim model) and preventing the projection head from training properly.

---

### Potential Contributions & Next Steps

I am happy to put together a PR (or split into two PRs) to address this and modernize the integration:

1. **Bug Fix PR:**
* Remove the $1/\sqrt{d_{\text{model}}}$ divisor when `algorithm == "adamw"` (or gate it strictly to `algorithm == "lion"` if Lion support is added).
* Ensure `lm_head` runs at the intended base scalar learning rate.

2. **Dion3 / Modernization PR:**
* Update `_import_dion()` and `DionConfig` to support the latest Dion3 release (`NorDion2`), exposing options for fractional updates (`fraction=0.25`), `selection_scope`, and megabatching.
* Update `MatrixAwareOptimConfig` parameter categorization to support 3D+ tensors (such as local 1D convolutions in hybrid architectures) via `flatten=True` rather than asserting and exiting.

Please let me know if you would like me to open a PR for the LM head fix first, or if you prefer a broader PR updating the integration to Dion3!