(model, num_to_forward_prefetch)
| 14 | return has_gpu and gpu_count >= min_gpus |
| 15 | |
| 16 | def set_modules_to_forward_prefetch(model, num_to_forward_prefetch): |
| 17 | for i, layer in enumerate(model.layers): |
| 18 | if i >= len(model.layers) - num_to_forward_prefetch: |
| 19 | break |
| 20 | layers_to_prefetch = [ |
| 21 | model.layers[i + j] for j in range(1, num_to_forward_prefetch + 1) |
| 22 | ] |
| 23 | layer.set_modules_to_forward_prefetch(layers_to_prefetch) |
| 24 | |
| 25 | |
| 26 | def set_modules_to_backward_prefetch(model, num_to_backward_prefetch): |