[megatron] megatron remove num_samples

[Jintao-Huang]

No description provided.

[gemini-code-assist]

Code Review

This pull request simplifies batch preparation and data collation by removing the num_samples parameter and instead deriving it dynamically from the length of seq_lens when packed_seq_params is present. While this cleanup streamlines the code, the review highlights a critical regression in get_last_tokens where packed_seq_params.num_samples is still accessed. Additionally, the feedback suggests adding safety checks when popping seq_lens from the batch and when accessing it within packed_seq_params to prevent potential KeyError or AttributeError exceptions.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.