[QUESTION] Document indices extended in add_document function in IndexedDatasetBuilder but not in add_item

[clarence-lee-sheng]

Dear Megatron-LM team

In the IndexedDatasetBuilder class add_item function (

Megatron-LM/megatron/core/datasets/indexed_dataset.py

Lines 767 to 779 in 912be78

	def add_item(self, tensor: torch.Tensor, mode: int = 0) -> None:
	"""Add a single item to the dataset
	Args:
	tensor (torch.Tensor): The item to add to the data file
	mode (int, optional): The mode for the item. Defaults to 0.
	"""
	np_array = numpy.array(tensor.numpy(), dtype=self.dtype)
	self.data_file.write(np_array.tobytes(order="C"))
	self.sequence_lengths.append(np_array.size)
	if self.multimodal:
	self.sequence_modes.append(mode)

), the document indices is not extended, however for the add_document function, the document indices are extended (

Megatron-LM/megatron/core/datasets/indexed_dataset.py

Lines 781 to 798 in 912be78

	def add_document(
	self, tensor: torch.Tensor, lengths: List[int], modes: Optional[List[int]] = None
	) -> None:
	"""Add an entire document to the dataset
	Args:
	tensor (torch.Tensor): The document to add
	lengths (List[int]): The lengths of each item in the document
	modes (Optional[List[int]], optional): The modes for each item in the document. Defaults to None.
	"""
	np_array = numpy.array(tensor, dtype=self.dtype)
	self.data_file.write(np_array.tobytes(order="C"))
	self.sequence_lengths.extend(lengths)
	self.document_indices.append(len(self.sequence_lengths))
	if self.multimodal:
	self.sequence_modes.extend(modes if modes is not None else [0] * lengths)

). May I ask why is that the case and when do we use add_item?

Thank you!

[Advait251206]

The distinction comes down to what each method is intended to represent.

• add_item() appends a single sequence/item to the dataset.
• add_document() appends an entire document, which may consist of one or more sequences, and records where that document ends.

The document_indices array is used to track document boundaries, not individual sequences. That's why add_document() updates it:

self.document_indices.append(len(self.sequence_lengths))

This records the cumulative number of sequences after the document has been added, allowing the dataset loader to reconstruct document boundaries later.

In contrast, add_item() only updates:

• the raw data,
• sequence_lengths,
• and optionally sequence_modes.

It intentionally does not update document_indices, because a single item isn't necessarily intended to represent a complete document. Multiple calls to add_item() could conceptually belong to the same document, so automatically extending document_indices after every item would effectively treat every item as a separate document, which isn't always correct.

As for when add_item() is used, it's generally intended for incremental dataset construction where items are written one at a time. If your downstream processing depends on document boundaries (for example, document-level sampling or preserving document context), you'll typically want to use add_document() instead so that document_indices is populated correctly.

So the differing behavior appears to be intentional: document_indices tracks documents, while sequence_lengths tracks individual sequences.

---

If this answer helped or pointed you in the right direction, I'd appreciate it if you could mark it as the accepted answer so it's easier for others with the same issue to find.

Also, if you found my contribution useful, I'd appreciate it if you could check out my GitHub profile, follow me, and star any repositories you find interesting.

GitHub: https://github.com/Advait251206