(model, linear_replacement, include_modules=['c_fc', 'c_proj'], copy_weights=True)
| 84 | # Replaces all linear layers with linear_replacement |
| 85 | # TODO: add int8 support for other linear layers including attn and convnets |
| 86 | def replace_linear(model, linear_replacement, include_modules=['c_fc', 'c_proj'], copy_weights=True): |
| 87 | for name, module in model.named_children(): |
| 88 | if len(list(module.children())) > 0: |
| 89 | replace_linear(module, linear_replacement, include_modules, copy_weights) |
| 90 | |
| 91 | if isinstance(module, torch.nn.Linear) and name in include_modules: |
| 92 | old_module = model._modules[name] |
| 93 | model._modules[name] = linear_replacement( |
| 94 | module.in_features, |
| 95 | module.out_features, |
| 96 | module.bias is not None, |
| 97 | ) |
| 98 | if copy_weights: |
| 99 | model._modules[name].weight.data.copy_(old_module.weight.data) |
| 100 | if model._modules[name].bias is not None: |
| 101 | model._modules[name].bias.data.copy_(old_module.bias) |
| 102 | |
| 103 | return model |
| 104 | |
| 105 | def convert_int8_model_to_inference_mode(model): |
| 106 | for m in model.modules(): |
nothing calls this directly
no outgoing calls
no test coverage detected