grad_component

Base class for Autograd Components that can be called and backpropagated through.

Classes

GradComponent(*args, **kwargs)

A base class to define interfaces for an auto-grad component/operator.

class GradComponent(*args, **kwargs)

Bases: Component

A base class to define interfaces for an auto-grad component/operator.

Compared with Component, GradComponent defines three important interfaces: - forward: the forward pass of the function, returns a Parameter object that can be traced and backpropagated. - backward: the backward pass of the function, updates the gradients/prediction score backpropagated from a “loss” parameter. - set_backward_engine: set the backward engine(a form of generator) to the component, which is used to backpropagate the gradients using LLM.

The __call__ method will check if the component is in training mode, and call the forward method to return a Parameter object if it is in training mode, otherwise, it will call the call method to return the output such as “GeneratorOutput”, “RetrieverOutput”, etc.

backward_engine: BackwardEngine

set_backward_engine(backward_engine: BackwardEngine, *args, **kwargs)

call(*args, **kwargs)

async acall(*args, **kwargs)

Implement this for your async call.

forward(*args, **kwargs) → Parameter

Default forward method for training: 1. for all args and kwargs, if it is a Parameter object, it will be tracked as Predecessor. 2. Trace input_args and full_response in the parameter object. 3. Return the parameter object.

TODO: all Gradcomponent should not allow args but only kwargs. For now, just check if id is in kwargs.

backward(*args, **kwargs)