SageMakerTrainingSessionRuntime / Client / execute_training_operation

execute_training_operation

SageMakerTrainingSessionRuntime.Client.execute_training_operation(**kwargs)

Submits an asynchronous training operation to the specified training session. The service executes operations in sequence number order. The session must be in Running or Blocked state. Use GetTrainingOperationResult to poll for the result. Supported operation types: forward, forwardBackward, optimizerStep, saveWeights, saveState, and loadState.

See also: AWS API Documentation

Request Syntax

response = client.execute_training_operation(
    trainingSessionName='string',
    sequenceNumber=123,
    clientToken='string',
    operationConfig={
        'forward': {
            'lossFnName': 'cross_entropy'|'importance_sampling'|'ppo'|'cispo',
            'datumList': [
                {
                    'modelInput': {
                        'chunks': [
                            {
                                'tokens': [
                                    123,
                                ],
                                'type': 'encoded_text'
                            },
                        ]
                    },
                    'lossFnInputs': {
                        'string': {
                            'data': {...}|[...]|123|123.4|'string'|True|None,
                            'dtype': 'string',
                            'shape': [
                                123,
                            ]
                        }
                    }
                },
            ],
            'lossFnConfig': {...}|[...]|123|123.4|'string'|True|None
        },
        'forwardBackward': {
            'lossFnName': 'cross_entropy'|'importance_sampling'|'ppo'|'cispo',
            'datumList': [
                {
                    'modelInput': {
                        'chunks': [
                            {
                                'tokens': [
                                    123,
                                ],
                                'type': 'encoded_text'
                            },
                        ]
                    },
                    'lossFnInputs': {
                        'string': {
                            'data': {...}|[...]|123|123.4|'string'|True|None,
                            'dtype': 'string',
                            'shape': [
                                123,
                            ]
                        }
                    }
                },
            ],
            'lossFnConfig': {...}|[...]|123|123.4|'string'|True|None
        },
        'optimizerStep': {
            'adamParamsConfig': {
                'learningRate': 123.0,
                'beta1': 123.0,
                'beta2': 123.0,
                'eps': 123.0,
                'weightDecay': 123.0,
                'gradClipNorm': 123.0,
                'loraPlusLrRatio': 123.0
            }
        },
        'saveWeights': {
            'mergeWeights': True|False,
            'description': 'string'
        },
        'saveState': {
            'description': 'string'
        },
        'loadState': {
            'modelPackageArn': 'string',
            'loadOptimizer': True|False
        }
    }
)
Parameters:

  • trainingSessionName (string) –

    [REQUIRED]

    The name of the training session to submit the operation to.

  • sequenceNumber (integer) –

    [REQUIRED]

    The monotonically increasing sequence number that determines execution order. The training session processes operations strictly in sequence number order. If a gap exists, the session blocks until the missing operation arrives or times out. Must be at least 1.

  • clientToken (string) – A unique, case-sensitive identifier that you provide to ensure that the operation completes no more than one time. If you retry a request with the same client token and the same parameters, the service ignores the request but does not return an error.

  • operationConfig (dict) –

    [REQUIRED]

    The configuration that specifies the type and parameters of the training operation. Exactly one member of the union must be set.

    Note

    This is a Tagged Union structure. Only one of the following top level keys can be set: forward, forwardBackward, optimizerStep, saveWeights, saveState, loadState.

    • forward (dict) –

      The configuration for a forward-only pass through the model without gradient computation.

      • lossFnName (string) – [REQUIRED]

        The name of the loss function to evaluate during the forward pass.

      • datumList (list) – [REQUIRED]

        The list of input data to feed through the model. Must contain at least one datum, and each datum must have a non-empty modelInput with at least one chunk.

        • (dict) –

          A single training datum containing model input and optional loss function inputs.

          • modelInput (dict) – [REQUIRED]

            The tokenized input to feed to the model.

            • chunks (list) – [REQUIRED]

              The ordered list of encoded text chunks forming the input sequence.

              • (dict) –

                A chunk of pre-tokenized text represented as a list of token IDs.

                • tokens (list) – [REQUIRED]

                  The token IDs representing the encoded text.

                  • (integer) –

                • type (string) – [REQUIRED]

                  The type of the chunk encoding.

          • lossFnInputs (dict) –

            The tensor inputs for the loss function, keyed by name. Unsupported keys are rejected. Supported keys: target_tokens, weights, advantages, logprobs, ref_log_prob, temperature. Required keys depend on the loss function type: SFT ( cross_entropy) requires target_tokens and weights; RL ( ppo, cispo, importance_sampling) requires target_tokens, weights, advantages, and logprobs.

            • (string) –

              • (dict) –

                A tensor with data, data type, and shape metadata.

                • data (document) – [REQUIRED]

                  The raw tensor data as JSON. The format is determined by the dtype field: integer arrays for int64, floating-point arrays for float32.

                • dtype (string) – [REQUIRED]

                  The data type of the tensor elements, such as int64 or float32.

                • shape (list) –

                  The dimensions of the tensor, such as [batch_size, seq_length].

                  • (integer) –

      • lossFnConfig (document) –

        Additional loss function configuration passed to the engine. Supported keys are clip_low_threshold (valid range: 0 to 1) and clip_high_threshold (valid range: 1 to 10). Unsupported keys are rejected.

    • forwardBackward (dict) –

      The configuration for a forward and backward pass that computes gradients.

      • lossFnName (string) – [REQUIRED]

        The name of the loss function to evaluate during the forward-backward pass.

      • datumList (list) – [REQUIRED]

        The list of input data to feed through the model. Must contain at least one datum, and each datum must have a non-empty modelInput with at least one chunk.

        • (dict) –

          A single training datum containing model input and optional loss function inputs.

          • modelInput (dict) – [REQUIRED]

            The tokenized input to feed to the model.

            • chunks (list) – [REQUIRED]

              The ordered list of encoded text chunks forming the input sequence.

              • (dict) –

                A chunk of pre-tokenized text represented as a list of token IDs.

                • tokens (list) – [REQUIRED]

                  The token IDs representing the encoded text.

                  • (integer) –

                • type (string) – [REQUIRED]

                  The type of the chunk encoding.

          • lossFnInputs (dict) –

            The tensor inputs for the loss function, keyed by name. Unsupported keys are rejected. Supported keys: target_tokens, weights, advantages, logprobs, ref_log_prob, temperature. Required keys depend on the loss function type: SFT ( cross_entropy) requires target_tokens and weights; RL ( ppo, cispo, importance_sampling) requires target_tokens, weights, advantages, and logprobs.

            • (string) –

              • (dict) –

                A tensor with data, data type, and shape metadata.

                • data (document) – [REQUIRED]

                  The raw tensor data as JSON. The format is determined by the dtype field: integer arrays for int64, floating-point arrays for float32.

                • dtype (string) – [REQUIRED]

                  The data type of the tensor elements, such as int64 or float32.

                • shape (list) –

                  The dimensions of the tensor, such as [batch_size, seq_length].

                  • (integer) –

      • lossFnConfig (document) –

        Additional loss function configuration passed to the engine. Supported keys are clip_low_threshold (valid range: 0 to 1) and clip_high_threshold (valid range: 1 to 10). Unsupported keys are rejected.

    • optimizerStep (dict) –

      The configuration for an optimizer step that updates model weights using accumulated gradients.

      • adamParamsConfig (dict) – [REQUIRED]

        The Adam optimizer hyperparameters for this step.

        • learningRate (float) – [REQUIRED]

          The learning rate for the optimizer. Must be greater than 0.

        • beta1 (float) –

          The exponential decay rate for the first moment estimates.

        • beta2 (float) –

          The exponential decay rate for the second moment estimates.

        • eps (float) –

          A small constant added to the denominator for numerical stability.

        • weightDecay (float) –

          The weight decay (L2 regularization) coefficient.

        • gradClipNorm (float) –

          The maximum gradient norm for gradient clipping. Gradients are scaled down if their norm exceeds this value.

        • loraPlusLrRatio (float) –

          The learning rate multiplier for LoRA+ B matrices relative to A matrices. When set, B matrices use learningRate * loraPlusLrRatio while A matrices use the base learning rate.

    • saveWeights (dict) –

      The configuration for saving model weights as a checkpoint.

      • mergeWeights (boolean) –

        Specifies whether to merge LoRA adapter weights into the base model before saving. When true, the saved checkpoint contains the full merged model weights.

      • description (string) –

        A human-readable description for the saved checkpoint.

    • saveState (dict) –

      The configuration for saving full training state, including model weights and optimizer state.

      • description (string) –

        A human-readable description for the saved state checkpoint.

    • loadState (dict) –

      The configuration for loading training state from a previously saved checkpoint.

      • modelPackageArn (string) – [REQUIRED]

        The ARN of the model package containing the checkpoint to load. Must be a valid SageMaker model package ARN in the format arn:aws:sagemaker:REGION:ACCOUNT:model-package/PACKAGE_NAME/VERSION. Model package group ARNs are not accepted.

      • loadOptimizer (boolean) – [REQUIRED]

        Specifies whether to also restore the optimizer state alongside model weights. When true, optimizer state is restored to resume training from where it left off.

Return type:

dict

Returns:

Response Syntax

{
    'operationId': 'string'
}

Response Structure

  • (dict) –

    • operationId (string) –

      The unique identifier for the submitted operation. Use this identifier with GetTrainingOperationResult to retrieve the operation status and result.

Exceptions