Fix VLM preprocessing and add mRoPE position handling in target head

[liusy58]

Motivation

This PR improves VLM training support in two areas:

• preserve dataset-provided system prompts during VLM conversation preprocessing
• carry image tensors through data processing for downstream model usage
• add multimodal RoPE position index generation in the target head

Modifications

Related Issues

Accuracy Test

Benchmark & Profiling

Checklist

• Format your code according to the Code Formatting with Pre-Commit.
• Add unit tests as outlined in the Running Unit Tests.
• Update documentation / docstrings / example tutorials as needed, according to Writing Documentation.
• Provide throughput / latency benchmark results and accuracy evaluation results as needed, according to Benchmark and Profiling and Accuracy Results.
• For reviewers: If you haven't made any contributions to this PR and are only assisting with merging the main branch, please remove yourself as a co-author when merging the PR.
• Please feel free to join our Slack channel at https://sgl-fru7574.slack.com/archives/C09784E3EN6 to discuss your PR.

[gemini-code-assist]

Code Review

This pull request updates VLM conversation preprocessing to support custom system prompts and ensures vision-related data fields are preserved during processing. It also introduces a get_rope_index method for calculating multimodal rotary position embeddings. Feedback includes addressing a potential IndexError in conversation handling, optimizing token search efficiency, and refactoring duplicated logic in the position ID calculation.

[gemini-code-assist]

There's a potential IndexError here. If source (which is examples["conversations"][i]) is an empty list, accessing source[0] will raise an exception. The check for an empty source on line 224 happens after this access.

To prevent this, you should add a check to ensure source is not empty before accessing its first element.

Suggested change

	if source[0]["role"] == "system":
	if source and source[0]["role"] == "system":

[gemini-code-assist]

The current implementation for finding ed_image and ed_video uses an in check followed by .index(), which results in scanning the input_tokens list twice for each token type in each iteration. This can be inefficient for long sequences.

A more efficient and Pythonic approach is to use a try-except block to handle cases where a token is not found. This avoids the redundant scan.

Suggested change

	ed_image = input_tokens.index(image_token_id, st) if image_token_id in input_tokens[st:] and remain_images > 0 else len(input_tokens) + 1
	ed_video = input_tokens.index(video_token_id, st) if video_token_id in input_tokens[st:] and remain_videos > 0 else len(input_tokens) + 1
	ed_image = len(input_tokens) + 1
	if remain_images > 0:
	try:
	ed_image = input_tokens.index(image_token_id, st)
	except ValueError:
	pass
	ed_video = len(input_tokens) + 1
	if remain_videos > 0:
	try:
	ed_video = input_tokens.index(video_token_id, st)
	except ValueError:
	pass

[gemini-code-assist]

The logic in this block to handle trailing text is very similar to the logic for handling text between vision tokens (lines 163-166). Specifically, calculating st_idx and appending text position IDs is duplicated.

To improve maintainability and reduce redundancy, consider refactoring this repeated logic. You could potentially merge this trailing text handling into the main loop or use a helper function to encapsulate the logic for appending text positions.