Memory inefficiency when loading attention_mask, causing dataloader OOM with long context

[shensimeteor]

Is your feature request related to a problem? Please describe.

In GPTSFTChatDataset::collate_fn, it seems return a huge tensor of attention_mask (code:
https://github.com/NVIDIA/NeMo/blob/main/nemo/collections/nlp/data/language_modeling/megatron/gpt_sft_chat_dataset.py#L380-L381)

The tensor has shape is like (n, 1, seq, seq). n is global_batch_size divided by data_parallel_size (aka, micro_batch_size * gradient_accumulation_steps, if my understanding is correct), seq is max context length in the batch. Therefore the tensor will be very huge when context length is long and when we have large gradient_accumulation_steps.

There're lots of wasted memory of the tensor: all subtensors of shape (1, seq, seq) are the same. So it's possible to reduce this tensor size by 1/n. It will help a lot with long context training.

Describe the solution you'd like

I'm not very familiar with the codebase so my thought could be wrong here. In GPTSFTChatDataset::collate_fn, we can just use a tensor of shape (1, seq, seq) to save the attention_mask. Then we need to modify or extend the get_iterator_k_split method used in GPTSFTModel (https://github.com/NVIDIA/NeMo-Aligner/blob/main/nemo_aligner/models/nlp/gpt/gpt_sft_model.py#L86). In the modified or extended version, instead of splitting attention_mask, we just create the actual (micro_batch, 1, seq, seq) of the attention_mask.

Describe alternatives you've considered

No

Additional context

No

[iwiwi]

I encountered the same issue, and the following workaround worked: adding get_attention_mask_from_fusion=True in the arguments given to dataset_cls in nemo_aligner/data/build_sft_dataset.py. By specifying this option, the attention mask matrices are instantiated directly in GPUs.

NeMo-Aligner/nemo_aligner/data/nlp/builders.py

Lines 421 to 451 in a2286e7

	dataset = dataset_cls(
	file_path=data_cfg.file_path,
	tokenizer=tokenizer,
	max_seq_length=data_cfg.max_seq_length,
	min_seq_length=data_cfg.min_seq_length,
	add_bos=data_cfg.get("add_bos", False),
	add_eos=data_cfg.get("add_eos", True),
	add_sep=data_cfg.get("add_sep", False),
	sep_id=0,
	max_num_samples=num_samples,
	seed=data_cfg.get("seed", 1234),
	label_key=data_cfg.get("label_key", "answer"),
	answer_only_loss=answer_only_loss,
	truncation_field=data_cfg.get("truncation_field", "text"),
	pad_to_max_length=data_cfg.get("pad_to_max_length", False),
	index_mapping_dir=data_cfg.get("index_mapping_dir", None),
	prompt_template=data_cfg.get("prompt_template", None),
	virtual_tokens=0,
	memmap_workers=data_cfg.get(
	"memmap_workers", None
	), # used to set num. of workers to create the memmap index files
	hf_dataset=data_cfg.get(
	"hf_dataset", False
	), # Whether to load the json file with the HuggingFace dataset. otherwise, will load the jsonl file with the JSONLMemMapDataset.
	truncation_method=data_cfg.get(
	"truncation_method", "right"
	), # used to choose truncation method. Options: ['random', 'left', 'right']
	special_tokens=special_tokens,
	output_original_text=data_cfg.get("output_original_text", False),
	**dataset_kwargs,
	)

    dataset = dataset_cls(
        ...
        output_original_text=data_cfg.get("output_original_text", False),
        get_attention_mask_from_fusion=True,  # Add this line
        **dataset_kwargs,
    )

Hope this helps @shensimeteor.

I also suggest that the NeMo-Aligner team add a get_attention_mask_from_fusion option like the following:

    dataset = dataset_cls(
        ...
        output_original_text=data_cfg.get("output_original_text", False),
        get_attention_mask_from_fusion=data_cfg.get("get_attention_mask_from_fusion", False),  # Add this line
        **dataset_kwargs,
    )