group_by_length silently does nothing for lazily-tokenized datasets

[kabirvashisht4-glitch]

Describe the bug

LengthGroupedSampler._compute_lengths has a "fast path" that unwraps .dataset
attributes until it finds a plain list, then indexes that list instead of the
dataset itself:

https://github.com/NVIDIA-NeMo/Automodel/blob/main/nemo_automodel/components/datasets/llm/length_grouped_sampler.py#L137-L153

# Fast path: access underlying list directly if available
raw = dataset
while hasattr(raw, "dataset"):
    raw = raw.dataset
if not isinstance(raw, list):
    raw = None
...
sample = raw[i] if raw is not None else dataset[i]
ids = sample.get("input_ids")
if ids is not None:
    lengths[i] = len(ids) if isinstance(ids, list) else ids.numel()

The LLM datasets in this repo tokenize lazily in __getitem__ and keep the
raw, untokenized rows in self.dataset. ChatDataset is the clearest case:
self.dataset is the list returned by _load_openai_messages (a plain
List[Dict] for local JSON/JSONL input), and input_ids only exists after
__getitem__ runs format_chat_template.

So the unwrap lands on rows shaped like {"messages": [...]},
sample.get("input_ids") returns None, and every length stays at the 0
initializer. sorted() on all-equal keys is stable, so sorted_indices is
just range(len(dataset)) — the sampler degrades to chunk-shuffled original
order and does no length grouping at all. There is no error and no warning.

The same unwrap is also unsafe for any wrapper that remaps indices
(e.g. torch.utils.data.Subset): raw[i] is not dataset[i], so lengths get
attributed to the wrong samples, and it can raise IndexError when
len(raw) < len(dataset).

Note the fast path buys nothing in the case it is actually correct: when
dataset is itself a plain list, the loop does not unwrap anything and
raw[i] is literally dataset[i]. It only changes behaviour in exactly the
cases where it is wrong.

Steps/Code to reproduce bug

group_by_length: true in the dataloader config with any lazily-tokenizing
dataset, e.g.:

dataset:
  _target_: nemo_automodel.components.datasets.llm.chat_dataset.ChatDataset
  path_or_dataset_id: /path/to/train.jsonl

dataloader:
  group_by_length: true

Minimal standalone repro (no tokenizer needed — same object shape as
ChatDataset: raw rows in .dataset, tokenization in __getitem__):

from nemo_automodel.components.datasets.llm.length_grouped_sampler import LengthGroupedSampler


class FakeChatDataset:
    def __init__(self, raw_rows):
        self.dataset = raw_rows  # raw, untokenized

    def __len__(self):
        return len(self.dataset)

    def __getitem__(self, idx):
        n = self.dataset[idx]["n_tokens"]
        return {"input_ids": list(range(n)), "labels": list(range(n))}


ds = FakeChatDataset([{"n_tokens": n} for n in [8, 128, 16, 64, 4, 256, 32, 512]])
sampler = LengthGroupedSampler(ds, batch_size=2, seed=0, num_replicas=1, rank=0)

print("computed lengths:", sampler.lengths)
print("actual lengths  :", [len(ds[i]["input_ids"]) for i in range(len(ds))])
print("sorted_indices  :", sampler.sorted_indices)

Output:

computed lengths: [0, 0, 0, 0, 0, 0, 0, 0]
actual lengths  : [8, 128, 16, 64, 4, 256, 32, 512]
sorted_indices  : [0, 1, 2, 3, 4, 5, 6, 7]

Batching that order at batch_size=2 costs 900 padding tokens; correct
length grouping costs 340.

Expected behavior

group_by_length: true groups similar-length samples so batches waste less
padding. Lengths should be read through dataset[i] whenever the unwrapped
list is not 1:1 with the dataset or does not already carry input_ids, and the
sampler should say something when it cannot determine any lengths instead of
silently becoming a no-op.

Environment overview

• Reproduced on main (0d1b8ce), CPU only — no GPU or distributed setup needed.

Additional context

Happy to send a PR: restrict the fast path to the case where it is provably
equivalent (unwrapped list is a list, same length as the dataset, and its
first row already has input_ids), otherwise go through dataset[i]; plus a
warning when every computed length is zero.