(self)
| 79 | # Monkeypatch this so it counts all layer parameters, not just trainable parameters. |
| 80 | # This helps it divide the layers between GPUs more evenly when training a LoRA. |
| 81 | def _count_all_layer_params(self): |
| 82 | param_counts = [0] * len(self._layer_specs) |
| 83 | for idx, layer in enumerate(self._layer_specs): |
| 84 | if isinstance(layer, ds_pipe_module.LayerSpec): |
| 85 | l = layer.build() |
| 86 | param_counts[idx] = sum(p.numel() for p in l.parameters()) |
| 87 | elif isinstance(layer, nn.Module): |
| 88 | param_counts[idx] = sum(p.numel() for p in layer.parameters()) |
| 89 | return param_counts |
| 90 | ds_pipe_module.PipelineModule._count_layer_params = _count_all_layer_params |
| 91 | |
| 92 |
nothing calls this directly
no outgoing calls
no test coverage detected