default_collater reshapes the caller's example tensors from [S] to [1, S]

Describe the bug

default_collater builds pre-batched tensor fields with:

ans[key] = torch.cat([batchify(v) for v in values], dim=0)

values are the tensors taken straight out of the caller's examples, and
batchify unsqueezes a 1-D tensor in place:

if tensor.ndim == 1:
    return tensor.unsqueeze_(0)

So collating a batch reshapes each example's own tensor from [S] to [1, S].

To be clear about where the defect is: batchify's in-place behavior looks deliberate
— test_batchify_adds_batch_dimension pins it with assert out is vec. The bug is that
default_collater applies that in-place helper to tensors it does not own. Every other
batchify call site passes a freshly built tensor (torch.stack(...),
torch.LongTensor(...)), so this is the only site where the mutation escapes.

Steps/Code to reproduce bug

import torch
from nemo_automodel.components.datasets.utils import default_collater

sample = {"input_ids": torch.arange(4), "labels": torch.arange(4)}
other  = {"input_ids": torch.arange(4), "labels": torch.arange(4)}

print(tuple(sample["input_ids"].shape))   # (4,)
out = default_collater([sample, other])
print(tuple(out["input_ids"].shape))      # (2, 4)   <- batch is fine
print(tuple(sample["input_ids"].shape))   # (1, 4)   <- example was rewritten
print(len(sample["input_ids"]))           # 1        <- was 4

Expected behavior

Collating reads the examples and returns a batch. The examples are unchanged
afterwards.

Actual behavior

Each example's tensor field is left reshaped to [1, S]. The collated batch itself is
correct, and re-collating the same examples still yields the right [B, S], so nothing
raises — the damage is to the example objects, which the dataset may hand out again.

len(sample["input_ids"]) becomes 1, and sample["input_ids"].shape[0] becomes 1
instead of S. Any later reader of that object sees the wrong length: a second epoch
over an in-memory dataset, a metric computed off the sample, or packing code that
measures ids.shape[0] (neat_pack_dataset does exactly that).

Scope

This is latent today rather than actively corrupting training runs. The collated batch
is correct, and the one in-tree consumer that measures a sample after construction,
LengthGroupedSampler._compute_lengths, happens to use ids.numel(), which is
shape-independent. It bites whoever next reads a collated example by length or shape.

Existing coverage cannot catch it: no test asserts anything about the input examples
after default_collater returns.

Environment overview

• Reproduced on CPU, torch 2.x, current main (23610d979). No GPU required.

Proposed fix

Add the batch axis out-of-place at the call site and leave batchify's tested contract
alone, plus a docstring warning on batchify so the next caller sees the side effect.
PR to follow.

Separately

pad_within_micro(batch, pad_token_id=None) pads every row with batch[0][-1] —
the first example's last token — rather than each row's own:

pad_within_micro([[10, 11, 12, 13, 14], [20, 21, 22]], None)
# -> [[10, 11, 12, 13, 14], [20, 21, 22, 14, 14]]
#                                      ^^  ^^ from example 0

get_pad_token_from_key returns None for any field outside
{labels, attention_mask, loss_mask, input_ids} — position_ids among them — so those
fields get one example's trailing value copied into another's padding. Filing it here as
an observation rather than folding it into the same PR, since the intended semantics of
the None branch aren't obvious from the code. Happy to split it out if that's useful.