[Training] The gradient builder has not been registered for node with op type MatMulNBits

[Positronx]

Describe the issue

I noticed that the gradient builder has not been registered for a multitude of ONNX operators. For example, is there a plan to implement the gradient operator for the node with op type MatMulNBits in the future?

To reproduce

Try to generate training artifacts for any ONNX model containing the node MatMulNbis

Urgency

No response

ONNX Runtime Installation

Built from Source

ONNX Runtime Version or Commit ID

1.18.0

PyTorch Version

2.2.0

Execution Provider

Default CPU

Execution Provider Library Version

CPU

[kshama-msft]

Thank you for reaching out with your question. Can you please specify an urgency and the scenario where you are using ORT for training. It will also be helpful if you can provide the list of ops that your scenario needs. Missing gradient support for certain ops is a known gap and the gradient builder will be implemented on a per-demand basis. If we understand your urgency, it will help us prioritize this work for implementation.

[Positronx]

Hi @kshama-msft,
Scenario : I want to finetune a quantized ONNX-genai model using ORTTraining so I'll need the gradient implementation of some specific nodes.

Urgency : as soon as possible

Here are the list of nodes that need a gradient implementation :
op_type : GroupQueryAttention, module : com.microsoft v1
op_type : SkipSimplifiedLayerNormalization, module : com.microsoft v1
op_type : SkipLayerNormalization, module : com.microsoft v1
op_type : MatMulNBits, module : com.microsoft v1 (I assume only 4 bit quantization is supported as of today)
op_type : MultiHeadAttention, module : com.microsoft v1 (If it isn't deprecated and replaced by GroupQueryAttention)

[kshama-msft]

Thank you for your response. Can you try exporting your model through torch.onnx.export and use that model as your starting point for training artifact generation? This exported model should have all ONNX ops and should temporarily unblock you as we work on gradient builder for the missing ops.

[skottmckay]

IIUC, typically a model would be trained at full precision (32-bit float values) and quantized after for more efficient inferencing. Why do you want/need to train the quantized model, and is there some evidence that doing so would be effective.

[Positronx]

Hi @skottmckay,
That's true for small models, but it's practically impossible to train large models on a single GPU using half presion (16-bit float values), let alone full precision since there isn't enough RAM to load all the model weights.
Edit : I don't recall that ONNXRuntime implemented a hybrid execution (CPU+GPU), because this will make it possible to train larger models.

[skottmckay]

AFAICS none of the training ops support quantized data types, so adding gradient implementations for the contrib ops would not be enough.

e.g.

onnxruntime/orttraining/orttraining/training_ops/cuda/nn/layer_norm.cc

Lines 32 to 37 in b1b8cb0

	REGISTER_GRADIENT_KERNEL_TYPED(float, float, float)
	REGISTER_GRADIENT_KERNEL_TYPED(double, double, double)
	REGISTER_GRADIENT_KERNEL_TYPED(MLFloat16, float, MLFloat16)
	REGISTER_GRADIENT_KERNEL_TYPED(float, float, MLFloat16)
	REGISTER_GRADIENT_KERNEL_TYPED(MLFloat16, float, float)
	REGISTER_GRADIENT_KERNEL_TYPED(BFloat16, float, BFloat16)

We're currently looking at making phi-3 work with the on-device training infrastructure so that just the last layer/s of the model could be trained with a much smaller memory requirement. But that would be with fp32 data using CPU.

[jingyanwangms]

@Positronx can you try instructions here to use ORTModule. https://github.com/microsoft/onnxruntime/blob/main/docs/ORTModule_Training_Guidelines.md
You just need to wrap the model with ORTModule.
I think at least GroupQueryAttention is not implemented in training. Make sure you use one of the supported types: float32/float16/bf16.

[Positronx]

@skottmckay, I see. Why wouldn't it possible to train phi-3 with fp16 using GPU?

[skottmckay]

fp16 is fine if you have GPU.

if you're shipping a product to general users you can't assume there's a GPU though (and even if there is there are different libraries required for CUDA vs AMD vs integrated) - which is why we're looking at CPU based usage.

[github-actions]

This issue has been automatically marked as stale due to inactivity and will be closed in 30 days if no further activity occurs. If further support is needed, please provide an update and/or more details.