**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 ```python # 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.: ```yaml 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__`): ```python 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` (0d1b8ce9), 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.