Fix CrossEntropyLoss block to support multi-output models#28232
Open
Rishi-Dave wants to merge 1 commit intomicrosoft:mainfrom
Open
Fix CrossEntropyLoss block to support multi-output models#28232Rishi-Dave wants to merge 1 commit intomicrosoft:mainfrom
Rishi-Dave wants to merge 1 commit intomicrosoft:mainfrom
Conversation
…put models CrossEntropyLoss.build() created a SoftmaxCrossEntropyLoss node with two outputs (loss, log_prob) but never registered log_prob in model.graph.value_info. Graph optimizers then dropped the output def, causing the gradient builder to hit a C++ assertion (i < node_->OutputDefs().size()) via O(1) when generating training artifacts for models with multi-dimensional outputs (e.g. seq2seq). Fix: after appending the node, add a value_info entry for log_prob_output_name using the same elem_type as the input scores tensor. A guard prevents duplicate entries if build() is called more than once. This keeps the output def alive through graph cleanup without changing the user-visible API (the block still returns only loss_node_output_name). Fixes microsoft#22465
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
artifacts.generate_artifacts(..., loss=LossType.CrossEntropyLoss)no longer aborts withi < node_->OutputDefs().size()when the base model has multi-dimensional outputs.SoftmaxCrossEntropyLossop produces two outputs (loss,log_prob); the second was being dropped by graph optimizers because it had novalue_infoentry, leaving the gradient builder to dereference a missing output def viaO(1).Motivation
Fixes #22465. Users hit a hard C++ assertion when training models like DistilBERT whose forward graph emits a multi-dimensional last-hidden-state tensor. The same pattern appears for any seq2seq / LM training setup that pipes a 3-D output into
CrossEntropyLoss.This is a Python-only change scoped to the
onnxblocktraining-artifacts API; the core inference engine is unaffected.Changes
orttraining/orttraining/python/training/onnxblock/loss/loss.py— after appending theSoftmaxCrossEntropyLossnode, register avalue_infoentry forlog_prob_output_nameso its output def survives shape inference and graph cleanup. Idempotent — guarded against duplicate entries.orttraining/orttraining/test/python/orttraining_test_ort_apis_onnxblock.py— newtest_crossentropy_loss_multi_output_modelbuilds a 3-D output toy model, callsgenerate_artifactswithLossType.CrossEntropyLoss, and asserts the savedtraining_model.onnxretains both outputs on the SCE node.Test Plan
python -m pytest orttraining/orttraining/test/python/orttraining_test_ort_apis_onnxblock.py::test_crossentropy_loss_multi_output_model -vpython -m pytest orttraining/orttraining/test/python/orttraining_test_ort_apis_onnxblock.py -k crossentropy -vlintrunnerclean on the diff.Fixes #22465