[WIP] [megatron] support megatron seq_cls task_type #5759
Conversation
There was a problem hiding this comment.
Summary of Changes
Hello @Jintao-Huang, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!
This pull request lays the groundwork for integrating sequence classification capabilities into the Megatron framework. It introduces necessary configuration arguments, adapts the core GPT model to handle sequence classification outputs, and prepares the training pipeline for task-specific loss computation. The changes streamline problem type handling and ensure proper saving of task-specific model components.
Highlights
- Support for Sequence Classification (SeqCls): Introduces the
seq_clstask type within the Megatron framework, enabling models to perform sequence classification. - New Megatron Arguments: Adds
task_type,num_labels, andproblem_typetoExtraMegatronArgumentsto configure sequence classification tasks. - Custom Output Layer for SeqCls: Implements a dedicated
ColumnParallelLinearoutput layer in the Megatron GPT model specifically for sequence classification, usingnum_labelsfor output dimension. - Refactored Problem Type Handling: Simplifies the determination of
problem_typeby removing the dynamic_get_problem_typemethod and relying directly on the configuredproblem_typefrom model arguments. - Trainer Adaptation for SeqCls Loss: Modifies the trainer to conditionally use a specific
seq_cls_loss_funcfor sequence classification tasks, allowing for tailored loss computation. - Output Layer Saving: Ensures the newly introduced
output_layerfor sequence classification is correctly included in the modules to be saved during training.
Using Gemini Code Assist
The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.
Invoking Gemini
You can request assistance from Gemini at any point in your pull request via creating an issue comment (i.e. comment on the pull request page) using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands.
| Feature | Command | Description |
|---|---|---|
| Code Review | /gemini review |
Performs a code review for the current pull request in its current state. |
| Pull Request Summary | /gemini summary |
Provides a summary of the current pull request in its current state. |
| Comment | @gemini-code-assist | Responds in comments when explicitly tagged, both in issue comments and review comments. |
| Help | /gemini help |
Displays a list of available commands. |
Customization
To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.
Limitations & Feedback
Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.
You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.
Footnotes
-
Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩
There was a problem hiding this comment.
Code Review
This pull request adds support for sequence classification tasks (seq_cls) within the Megatron framework. The changes include adding necessary arguments, modifying the GPT model to include a classification head, and updating the trainer to handle a new loss function. The logic for determining problem_type has been simplified, now relying on explicit configuration rather than inference. My review focuses on the trainer implementation, where the sequence classification loss function is currently a placeholder and is not called correctly. These are critical issues that need to be addressed to make the feature functional.
| def seq_cls_loss_func(self, output_tensor, *, labels: torch.Tensor, packed_seq_params=None): | ||
| pass |
There was a problem hiding this comment.
The seq_cls_loss_func is not implemented. It should compute the loss for sequence classification tasks based on the problem_type. Note that megatron's train_step expects the loss function to return a tuple of (unreduced_loss, num_tokens, loss_reduced_dict). The dictionary in loss_reduced_dict is used for logging and its key is expected to be 'lm loss' in the default megatron code.
def seq_cls_loss_func(self, output_tensor, *, labels: torch.Tensor, packed_seq_params=None):
args = get_args()
logits = output_tensor
if packed_seq_params:
# padding_free
indices = packed_seq_params.cu_seqlens_q[1:] - 1
logits = logits[0, indices]
else:
# This may not be correct if there is padding.
# It's better to use padding_free for seq_cls task.
logger.warning('It is recommended to use `padding_free` for `seq_cls` task.')
logits = logits[:, -1]
if args.problem_type == 'single_label_classification':
loss_fct = torch.nn.CrossEntropyLoss()
loss = loss_fct(logits.view(-1, args.num_labels), labels.view(-1))
elif args.problem_type == 'multi_label_classification':
loss_fct = torch.nn.BCEWithLogitsLoss()
loss = loss_fct(logits, labels.to(logits.dtype))
elif args.problem_type == 'regression':
loss_fct = torch.nn.MSELoss()
loss = loss_fct(logits.squeeze(), labels.squeeze().to(logits.dtype))
else:
raise ValueError(f'Unsupported problem_type: {args.problem_type}')
num_samples = torch.tensor(
labels.size(0), device=labels.device, dtype=torch.int)
loss_reduced = loss.clone().detach()
if mpu.get_data_parallel_world_size() > 1:
torch.distributed.all_reduce(loss_reduced, group=mpu.get_data_parallel_group())
loss_reduced /= mpu.get_data_parallel_world_size()
# The key 'lm loss' is expected by megatron for logging.
return (
loss,
num_samples,
{
'lm loss': (loss_reduced, num_samples)
},
)| channels=channels, | ||
| packed_seq_params=packed_seq_params) | ||
| if self.args.task_type == 'seq_cls': | ||
| loss_func = self.seq_cls_loss_func(output_tensor, labels=labels, packed_seq_params=packed_seq_params) |
There was a problem hiding this comment.
The seq_cls_loss_func is being called directly, but forward_step should return a callable loss_func. You should wrap the call in functools.partial to create a callable, similar to the else branch.
loss_func = partial(self.seq_cls_loss_func, labels=labels, packed_seq_params=packed_seq_params)
No description provided.