megatron.text_generation.forward_step.ForwardStep

class megatron.text_generation.forward_step.ForwardStep(model, max_batch_size, max_sequence_len)

Bases: object

Forward step function with all the communications. We use a class here to hide the inference parameters from the outside caller.

__call__(tokens, position_ids, attention_mask)

Invocation of the forward methods. Note that self.inference_params is being modified by the forward step.