def _sync_params()

in summarize_from_feedback/query_response_model.py [0:0]

• 19 lines of code
• 6 McCabe index (conditional complexity)

    def _sync_params(self, params_to_init, heads_to_init):
        if self.layout.n_replicas > 1:
            for param in params_to_init:
                self.dp_comm.broadcast(
                    param.data,
                    src=self.layout.dp_sibling_ranks[0],
                    name="broadcast_params_from_zeroeth_replica",
                )

        if self.layout.n_shards > 0:
            params_to_sync_shards = []
            for head in heads_to_init:
                params_to_sync_shards.append(self.model.scalar_heads[head].weight)
                params_to_sync_shards.append(self.model.scalar_heads[head].bias)
            for param in params_to_sync_shards:
                self.model.mp_comm.broadcast(
                    param.data,
                    src=self.layout.mp_sibling_ranks[0],
                    name="broadcast_params_from_zeroeth_shard",
                )