Distillation loss path upcasts the full [B, S, V] logits to fp32 when only K columns are read

*(Edited: this originally led with the `logprob_chunk_size` point, now the second section. While checking my shapes against the shipped configs I found something simpler and more general, so I reordered and corrected the framing.)*

## The full logits are upcast to fp32 before it is known that only K columns are needed

`get_distillation_topk_logprobs_from_logits` opens with

```python
student_logits = student_logits.to(torch.float32) # [B, S, V]
```

Three of its four branches read only `K` columns from that tensor:

- `zero_outside_topk=False`, non-TP → `student_logits.gather(-1, teacher_topk_indices)`
- `zero_outside_topk=False`, TP → `gather_logits_at_global_indices(...)`, which casts each chunk to fp32 internally
- `zero_outside_topk=True`, TP → `ChunkedDistributedGatherLogprob`, likewise

Only `zero_outside_topk=True` on the non-TP path requires full-vocab fp32, for its `log_softmax`.

Gathering before the cast is equivalent, since `gather` is a selection and `to()` is elementwise, and produces bitwise-identical values and gradients. Measured on an RTX 4090 at B=1, V=151,936, K=64, the values shipped in `distillation_math.yaml` (`train_micro_batch_size: 1`, `topk_logits_k: 64`, `Qwen3-1.7B-Base`):

| seq_len | cast-then-gather (current) | gather-then-cast | saved |
|---:|---:|---:|---:|
| 2048 | 1.74 GiB | 0.58 GiB | 1.16 GiB |
| 8192 (config default) | 6.96 GiB | 2.33 GiB | 4.63 GiB |

At seq 8192 the gather also drops from 9.2 ms to 0.1 ms, since the current form moves 2.3 GiB through memory to retain 64 columns per position. This is an improvement on both axes rather than a tradeoff, and `zero_outside_topk=False` is the default, so it applies to the common path.

Filed as #3496.

## Separate point: `logprob_chunk_size` does not reach these paths

`policy.logprob_chunk_size` is threaded through the logprob paths (`automodel/train.py:733,744`, `dtensor_policy_worker.py:1224,1234`, `megatron/train.py:596`, `loss/utils.py:93,354`) but not through these two:

1. `get_distillation_topk_logprobs_from_logits` accepts no `chunk_size`; its TP branch hardcodes `max(1, min(seq_len_local, 1024))` instead of the configured value.
2. `student_next_token_ce` (`x_token/loss_utils.py:347`) calls `get_logprobs_from_vocab_parallel_logits(logits, input_ids, seq_index=seq_index)` without one.

A correction to my original framing: the figures I first quoted here, 27.83 GiB down to 7.98 GiB at seq 8192, require `zero_outside_topk=true`, which is not the default. `distillation_math.yaml` sets it `false` and the schema default is `False`. That part is therefore narrower than I presented it. It is also a genuine tradeoff of roughly 1.4x wall-clock, since the backward recomputes the per-chunk softmax.

Unlike the first section this is a judgement call, and I would want your read on it, primarily on what share of distillation users run DTensor rather than Megatron, given that the TP branch already chunks.