# Model

## Save/Load

Ref: <https://pytorch.org/tutorials/beginner/saving\\_loading\\_models.html>

* `torch.save(model.state_dict(), PATH)`
  * only save state\_dict, which is an OrderedDict containing trained weights for each layer
  * use `model = TheModelClass(*args, **kwargs)load_state_dict(torch.load(PATH))` to load
* `torch.save(model, PATH)`
  * save entire things about Model using **`pickle`**, since it will serialize all related things, it has a more restricted environment to load, including the defined model, the dictionary structure and so on.

## Functions

### [no\_sync](https://pytorch.org/docs/stable/generated/torch.nn.parallel.DistributedDataParallel.html#torch.nn.parallel.DistributedDataParallel.no_sync)

> A context manager to disable gradient synchronizations across DDP processes. Within this context, gradients will be accumulated on module variables, which will later be synchronized in the first forward-backward pass exiting the context.

Example in Megatron [forward\_backward\_no\_pipelining](https://github.com/NVIDIA/Megatron-LM/blob/b44dca25727c294a7f825e74a3c4a53744cc8404/megatron/schedules.py#L223):

```py
context_handler = dummy_handler
if isinstance(model, torchDDP):
    context_handler = model.no_sync

losses_reduced = []
input_tensor, output_tensor_grad = None, None
with context_handler():
    for i in range(get_num_microbatches() - 1):
        output_tensor = forward_step(forward_step_func, data_iterator, model,
                                      input_tensor, losses_reduced)
        if not forward_only:
            backward_step(optimizer, input_tensor, output_tensor,
                          output_tensor_grad)

# Run computation for last micro-batch out of context handler (want to
# synchronize gradients).
output_tensor = forward_step(forward_step_func, data_iterator, model,
                              input_tensor, losses_reduced)
if not forward_only:
    backward_step(optimizer, input_tensor, output_tensor, output_tensor_grad)
```

## ONNX

### ONNX export

```python
torch.onnx.export(
    model,
    args,
    output_path,
    opset_version: int = None,
    input_names: List = None,
    output_names: List = None,
    dynamic_axes: dict = None,
    ...
)
```

* `args`: The support for kwargs is not sufficient, recommend to use positional arguments instead. `args` will be interpreted as (\*args) and passed to the model, therefore it should have a wrapper at the outmost. For example, if the model has several inputs, it should be like:

  ```python
  torch.onnx.export(model, (input1, input2, input3), ...)
  ```
* `input_names` & `output_names`: will be readable in onnx model if setted
* `dynamic_axes`: specify the dynamic axes of input(`input_names` required). Recommend to use dict inside dict, for example:

  ```python
  torch.onnx.export(model, (input1, input2, input3), 
    dynamic_axes={
        'input1': {
            0: 'batch',
            1: 'sequence_length'
        },
        'input2': {
            0: 'batch',
            2: 'hidden_state'
        },
        'input3': {
            1: 'sequence_length'
        }
    }
  )
  ```

  In this case, onnx model will replace the axes with its dynamic name in each input parameter
* `opset_version`: This specify the op set version used for generating onnx model. For example, dynamic slices requires opset>=10, therefore at least set opset=10 for normal execution.


---

# Agent Instructions: Querying This Documentation

If you need additional information that is not directly available in this page, you can query the documentation dynamically by asking a question.

Perform an HTTP GET request on the current page URL with the `ask` query parameter:

```
GET https://legacy.cookielau.com/framework/pytorch/01-model.md?ask=<question>
```

The question should be specific, self-contained, and written in natural language.
The response will contain a direct answer to the question and relevant excerpts and sources from the documentation.

Use this mechanism when the answer is not explicitly present in the current page, you need clarification or additional context, or you want to retrieve related documentation sections.