fix(datasets): stop default_collater reshaping the caller's example tensors

[kabirvashisht4-glitch]

What does this PR do ?

Fixes #3933: default_collater reshaped the caller's example tensors as a side effect of
collating.

The pre-batched branch built its output with torch.cat([batchify(v) for v in values]),
where values are the tensors taken straight out of the caller's examples. batchify
unsqueezes a 1-D tensor in place and returns the same object, so each example's field
was rewritten from [S] to [1, S]:

sample = {"input_ids": torch.arange(4)}
out = default_collater([sample, other])
out["input_ids"].shape      # (2, 4)   <- batch is correct
sample["input_ids"].shape   # (1, 4)   <- example was rewritten
len(sample["input_ids"])    # 1        <- was 4

Where the fix goes, and why not in batchify

batchify's in-place behavior is deliberate, not a typo — test_batchify_adds_batch_dimension
pins it with assert out is vec and a comment saying so. Changing unsqueeze_ to
unsqueeze would break that contract for every caller to fix one of them.

Every other batchify call site passes a freshly constructed tensor (torch.stack(...),
torch.LongTensor(...)), so default_collater's pre-batched branch is the only place the
mutation escapes into data the function does not own. The fix is therefore at that call
site: add the batch axis out-of-place. batchify keeps its contract and gains a docstring
warning so the next caller sees the side effect before hitting it.

Impact

Latent rather than actively corrupting runs, and I'd rather say so than oversell it. The
collated batch is correct, re-collating the same examples still yields the right [B, S],
and the one in-tree consumer that measures a sample after construction —
LengthGroupedSampler._compute_lengths — happens to use ids.numel(), which is
shape-independent.

What is broken is the example object the dataset may hand out again:
len(sample["input_ids"]) and sample["input_ids"].shape[0] both become 1. That bites a
second epoch over an in-memory dataset, a metric read off the sample, or packing code
measuring ids.shape[0] — neat_pack_dataset does exactly that.

Changelog

• Build the pre-batched field with an out-of-place v.unsqueeze(0) instead of
  batchify(v), so default_collater no longer mutates the examples it was given.
• Document on batchify that a 1-D input is unsqueezed in place and returned as the same
  object, and on default_collater that — apart from the existing ___PAD_TOKEN_IDS___
  pop — inputs are left unmodified.
• Add two CPU tests: one that the examples keep their [S] shape and len() across a
  collate, and one that collating the same example objects twice (an in-memory dataset
  across epochs) leaves them intact and gives the same batch.

Before your PR is "Ready for review"

Pre checks:

• Make sure you read and followed Contributor guidelines
• Did you write any new necessary tests?
• Did you add or update any necessary documentation?

Verified locally:

• Both new tests fail on main and pass here. Reverting only the nemo_automodel/ change
  and keeping the tests reproduces both failures.
• pytest tests/unit_tests/datasets/ -> 1361 passed, 24 failed. The same 24 fail on clean
  main (1359 passed there); they are pre-existing and unrelated (object-storage and
  retrieval-launcher tests). The delta is exactly the two tests added here.
• ruff format --check and ruff check clean on both changed files.
• CPU only, no GPU needed.

Additional Information

• I also noticed that pad_within_micro(batch, pad_token_id=None) pads every row with
  batch[0][-1] rather than each row's own last token, which affects fields with no
  default pad token such as position_ids. It is written up at the end of default_collater reshapes the caller's example tensors from [S] to [1, S] #3933 and
  deliberately left out of this PR, since the intended semantics of that branch are not
  obvious from the code. Happy to split it out into its own change.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[yuhezhang-ai]

/ok to test 7f69807