Initialize affine weight for model parallel on GPU.
(weight, init_method, partition_dim, stride=1)
| 81 | |
| 82 | |
| 83 | def _initialize_affine_weight_gpu(weight, init_method, partition_dim, stride=1): |
| 84 | """Initialize affine weight for model parallel on GPU.""" |
| 85 | |
| 86 | set_tensor_model_parallel_attributes( |
| 87 | tensor=weight, is_parallel=True, dim=partition_dim, stride=stride |
| 88 | ) |
| 89 | |
| 90 | if ds_checkpointing.is_configured(): |
| 91 | global get_cuda_rng_tracker |
| 92 | get_cuda_rng_tracker = ds_checkpointing.get_cuda_rng_tracker |
| 93 | |
| 94 | with get_cuda_rng_tracker().fork(): |
| 95 | init_method(weight) |
| 96 | |
| 97 | |
| 98 | def _initialize_affine_weight_cpu( |
no test coverage detected