Add files using upload-large-folder tool

b025706 verified 25 days ago

13.6 kB

# SPDX-FileCopyrightText: © 2024 Tenstorrent USA, Inc.

# SPDX-License-Identifier: Apache-2.0

import math

import torch

import ttnn
from models.common.lightweightmodule import LightweightModule
from models.common.utility_functions import copy_to_buffer
from models.tt_transformers.tt.ccl import tt_all_reduce
from models.tt_transformers.tt.common import Mode


def _lm_head_cache_file_name(
    weight_cache_path,
    *,
    dummy_weights,
    num_splits,
    shard,
    width,
    mode,
    galaxy_2d,
    mesh_shape,
):
    if dummy_weights:
        return None

    layout_suffix = f"_galaxy_2d_{mesh_shape[0]}x{mesh_shape[1]}_v1" if galaxy_2d else ""
    return weight_cache_path / (f"output_lm_head_{num_splits}_split_shard_{shard}_{width}_mode_{mode}{layout_suffix}")


class LMHead(LightweightModule):
    def __init__(
        self,
        args,
        mesh_device,
        tt_ccl,
        dtype,
        state_dict,
        state_dict_prefix,
        weight_cache_path,
        max_columns_per_device,  # too many columns per device lead to L1 OOM
        prefetcher=None,
    ):
        super().__init__()
        self.args = args
        self.mesh_device = mesh_device
        self.tt_ccl = tt_ccl
        self.dtype = dtype
        self.vocab_size = args.vocab_size
        self.padded_vocab_size = args.padded_vocab_size
        self.num_devices = args.num_devices
        self.prefetcher = prefetcher
        self.use_galaxy_no_prefetch = args.is_galaxy and prefetcher is None

        vocab_shards = args.cluster_shape[0] if self.use_galaxy_no_prefetch else self.num_devices
        size_per_device = self.padded_vocab_size // vocab_shards

        tile_size = 32
        max_columns_per_device_ring_mm = math.ceil((max_columns_per_device) / tile_size) * tile_size
        max_columns_per_device_dram_sharded = max_columns_per_device

        self.model_config = args.get_model_config()

        num_splits_ring_mm = math.ceil(size_per_device / max_columns_per_device_ring_mm)
        num_splits_dram_sharded = math.ceil(size_per_device / max_columns_per_device_dram_sharded)

        self.split_sizes_dram_sharded = [min(size_per_device, max_columns_per_device_dram_sharded)] * (
            num_splits_dram_sharded - 1
        )
        self.split_sizes_dram_sharded.append(size_per_device - sum(self.split_sizes_dram_sharded))  # remaining columns

        self.split_sizes_ring_mm = [min(size_per_device, max_columns_per_device_ring_mm)] * (num_splits_ring_mm - 1)
        self.split_sizes_ring_mm.append(size_per_device - sum(self.split_sizes_ring_mm))  # remaining columns

        # Split the output weights
        torch_output_weights = state_dict[f"{state_dict_prefix}output.weight"].permute(1, 0)

        # Pad the output weights to the padded vocab size with zeros
        if self.vocab_size < self.padded_vocab_size:
            padding_size = self.padded_vocab_size - self.vocab_size
            torch_output_weights = torch.cat(
                [
                    torch_output_weights,
                    torch.zeros(torch_output_weights.shape[0], padding_size, dtype=torch_output_weights.dtype),
                ],
                dim=-1,
            )

        self.output_weights_dram_sharded = []
        self.output_weights_ring_mm = []

        self.split_sizes = [self.split_sizes_dram_sharded]
        if self.prefetcher is not None:
            self.split_sizes.append(self.split_sizes_ring_mm)

        for mode, split_sizes in enumerate(self.split_sizes):
            for i, split_size in enumerate(split_sizes):
                # Create a list to store the split tensors for each device
                device_splits = []
                for device in range(vocab_shards):
                    start = device * size_per_device + sum(split_sizes[:i])
                    end = start + split_size
                    device_splits.append(torch_output_weights[:, start:end])

                # Concatenate the splits from all devices
                combined_split = torch.cat(device_splits, dim=-1)

                cache_file_name = _lm_head_cache_file_name(
                    weight_cache_path,
                    dummy_weights=args.dummy_weights,
                    num_splits=len(split_sizes),
                    shard=i,
                    width=combined_split.shape[-1],
                    mode=mode,
                    galaxy_2d=self.use_galaxy_no_prefetch and mode == 0,
                    mesh_shape=args.cluster_shape,
                )

                def pad_to_power_of_2(n):
                    if n <= 0:
                        return 1
                    return 1 << (n - 1).bit_length()

                if mode == 0:
                    if self.use_galaxy_no_prefetch:
                        torch_weight = combined_split.unsqueeze(0).unsqueeze(0)
                        mesh_mapper = ttnn.ShardTensor2dMesh(mesh_device, dims=(3, 2), mesh_shape=args.cluster_shape)
                        memory_config = ttnn.DRAM_MEMORY_CONFIG
                    else:
                        torch_weight = combined_split
                        mesh_mapper = ttnn.ShardTensorToMesh(mesh_device, dim=-1)
                        memory_config = args.create_dram_sharded_mem_config(
                            k=args.dim, n=math.ceil(combined_split.shape[-1] / self.num_devices)
                        )
                    self.output_weights_dram_sharded.append(
                        ttnn.as_tensor(
                            torch_weight,
                            device=mesh_device,
                            mesh_mapper=mesh_mapper,
                            layout=ttnn.TILE_LAYOUT,
                            dtype=dtype,
                            memory_config=memory_config,
                            cache_file_name=cache_file_name,
                        )
                    )
                else:
                    memory_config = args.create_dram_sharded_mem_config(
                        k=args.dim,
                        n=pad_to_power_of_2(math.ceil(combined_split.shape[-1] / self.num_devices)),
                        dram_grid=self.prefetcher.to_core_range_set(self.prefetcher.dram_banks()),
                    )
                    self.output_weights_ring_mm.append(
                        ttnn.as_tensor(
                            combined_split,
                            device=mesh_device,
                            mesh_mapper=ttnn.ShardTensorToMesh(mesh_device, dim=-1),
                            layout=ttnn.TILE_LAYOUT,
                            dtype=dtype,
                            memory_config=memory_config,
                            cache_file_name=cache_file_name,
                        )
                    )

        self.compute_kernel_config = ttnn.WormholeComputeKernelConfig(
            math_fidelity=ttnn.MathFidelity.HiFi2,
            math_approx_mode=False,
            fp32_dest_acc_en=False,
            packer_l1_acc=True,
        )

    def _update_output_weights_dram_sharded(self, weight: ttnn.Tensor) -> None:
        """In-place replace every chunk of ``self.output_weights_dram_sharded``
        on device (``weight`` is HF ``(1, 1, V, H)``, replicated, TILE, bf16,
        DRAM-interleaved). Mirrors the constructor's host-side pad+transpose+split
        but on device: optional ``ttnn.pad`` to ``padded_vocab_size``,
        ``ttnn.transpose``, contiguous per-chunk ``ttnn.slice``, then
        ``copy_to_buffer`` reshards into the DRAM-sharded dest (preserving
        addresses). Single-device only; multi-device per-device gather not
        implemented.
        Tile-alignment caveat: ``ttnn.slice`` on TILE requires offsets aligned
        to TILE_SIZE (32); the default Llama configs satisfy this, but assert it
        rather than silently miscompile slices.
        """
        assert self.num_devices == 1, (
            "LMHead.update for num_devices > 1 is not yet implemented; "
            "the multi-device path needs a per-device interleaved gather."
        )
        tile_size = 32
        assert all(s % tile_size == 0 for s in self.split_sizes_dram_sharded), (
            f"LMHead.update requires every split in self.split_sizes_dram_sharded "
            f"to be a multiple of TILE_SIZE={tile_size}, got "
            f"{self.split_sizes_dram_sharded}; on-device ttnn.slice cannot produce "
            "sub-tile-aligned chunks on TILE_LAYOUT."
        )

        padded = weight
        if self.vocab_size < self.padded_vocab_size:
            pad_amount = self.padded_vocab_size - self.vocab_size
            padded = ttnn.pad(padded, [(0, 0), (0, 0), (0, pad_amount), (0, 0)], value=0.0)

        permuted = ttnn.transpose(padded, -2, -1)

        start = 0
        for i, split_size in enumerate(self.split_sizes_dram_sharded):
            end = start + split_size
            chunk = ttnn.slice(permuted, (0, 0, 0, start), (1, 1, self.args.dim, end))
            copy_to_buffer(chunk, self.output_weights_dram_sharded[i], self.dtype)
            start = end

    def _update_output_weights_ring_mm(self, weight: ttnn.Tensor) -> None:
        """In-place replace every chunk of ``self.output_weights_ring_mm``. Not
        yet implemented: the ring-mm path (distinct per-chunk split, power-of-two
        column padding, prefetcher-derived DRAM grid) needs its own builder.
        """
        raise NotImplementedError("LMHead.update for output_weights_ring_mm (prefetcher path) is not yet implemented")

    def update(self, *, weight: ttnn.Tensor) -> None:
        """In-place replace the on-device LM-head weights via ``ttnn.copy``.
        HF-format input (see ``LLAMA_WEIGHT_TRANSFER.md``): ``weight`` is
        ``lm_head.weight``, shape ``(1, 1, vocab_size, hidden_size)``, bf16, TILE,
        DRAM-interleaved, replicated.
        Updates ``self.output_weights_dram_sharded`` on device (no host
        roundtrip). When the prefetcher is enabled the ring-mm mirror must also
        be synced; that path is not implemented. Every chunk keeps its device
        allocation, so captured traces and the prefetcher's recorded addresses
        stay valid.
        """
        self._update_output_weights_dram_sharded(weight)
        if len(self.output_weights_ring_mm) > 0:
            self._update_output_weights_ring_mm(weight)

    def forward(self, x: ttnn.Tensor, debug_input_torch=None, debug_weight_torch=None):
        outputs = []
        use_prefetcher = self.prefetcher is not None and self.prefetcher.mode == Mode.DECODE
        split_sizes = self.split_sizes_ring_mm if use_prefetcher else self.split_sizes_dram_sharded
        program_configs = [
            self.args.get_lm_head_program_config(split_size, self.prefetcher if use_prefetcher else None)
            for split_size in split_sizes
        ]

        output_weights = self.output_weights_ring_mm if use_prefetcher else self.output_weights_dram_sharded

        if self.use_galaxy_no_prefetch:
            x = ttnn.to_memory_config(x, ttnn.DRAM_MEMORY_CONFIG)

        self.lm_head_output_memory_config = self.args.get_lm_head_output_mem_config(
            Mode.DECODE if use_prefetcher else Mode.PREFILL, self.prefetcher if use_prefetcher else None
        )

        for i, (weight, pc) in enumerate(zip(output_weights, program_configs)):
            output = ttnn.linear(
                x,
                weight,
                compute_kernel_config=self.compute_kernel_config,
                program_config=None if self.use_galaxy_no_prefetch else pc,
                memory_config=(
                    ttnn.DRAM_MEMORY_CONFIG if self.use_galaxy_no_prefetch else self.lm_head_output_memory_config
                ),
                dtype=self.args.lm_head_dtype if hasattr(self.args, "lm_head_dtype") else ttnn.bfloat8_b,
                sub_device_id=self.prefetcher.worker_sub_device_id if use_prefetcher else None,
            )
            output = ttnn.to_memory_config(
                output,
                memory_config=self.args.get_lm_head_sharded_output_mem_config(
                    self.prefetcher if use_prefetcher else None
                ),
            )

            outputs.append(output)

        ttnn.deallocate(x)

        # Concatenate the outputs
        # outputs shape: a list of tensors, each tensor is 1,1,32,size_per_device per device
        # Decode traces retain these logits across requests. Honor the model's placement
        # so large vocabularies do not occupy L1 needed by the next eager prefill.
        output = ttnn.concat(
            outputs,
            dim=-1,
            memory_config=(
                self.args.model_config.get("LM_HEAD_OUTPUT_MEMCFG", ttnn.L1_MEMORY_CONFIG)
                if not use_prefetcher
                else ttnn.DRAM_MEMORY_CONFIG
            ),
            sub_core_grids=self.prefetcher.all_worker_cores_range_set if use_prefetcher else None,
        )

        # Only use reshard mem config for ring_mm mode
        if use_prefetcher:
            output = ttnn.to_memory_config(
                output,
                memory_config=self.args.get_lm_head_reshard_mem_config(self.prefetcher),
            )

        output = tt_all_reduce(
            output,
            self.mesh_device,
            self.tt_ccl,
            cluster_axis=1,
            dim=3 if self.args.is_galaxy else 0,
            memory_config=output.memory_config(),
            dtype=self.args.ccl_dtype,
            sharded=False,
            use_composite=True,
            subdevice_id=self.prefetcher.worker_sub_device_id if use_prefetcher else None,
        )

        return output

	# SPDX-FileCopyrightText: © 2024 Tenstorrent USA, Inc.

	# SPDX-License-Identifier: Apache-2.0

	import math

	import torch

	import ttnn
	from models.common.lightweightmodule import LightweightModule
	from models.common.utility_functions import copy_to_buffer
	from models.tt_transformers.tt.ccl import tt_all_reduce
	from models.tt_transformers.tt.common import Mode


	def _lm_head_cache_file_name(
	weight_cache_path,
	*,
	dummy_weights,
	num_splits,
	shard,
	width,
	mode,
	galaxy_2d,
	mesh_shape,
	):
	if dummy_weights:
	return None

	layout_suffix = f"_galaxy_2d_{mesh_shape[0]}x{mesh_shape[1]}_v1" if galaxy_2d else ""
	return weight_cache_path / (f"output_lm_head_{num_splits}_split_shard_{shard}_{width}_mode_{mode}{layout_suffix}")


	class LMHead(LightweightModule):
	def __init__(
	self,
	args,
	mesh_device,
	tt_ccl,
	dtype,
	state_dict,
	state_dict_prefix,
	weight_cache_path,
	max_columns_per_device, # too many columns per device lead to L1 OOM
	prefetcher=None,
	):
	super().__init__()
	self.args = args
	self.mesh_device = mesh_device
	self.tt_ccl = tt_ccl
	self.dtype = dtype
	self.vocab_size = args.vocab_size
	self.padded_vocab_size = args.padded_vocab_size
	self.num_devices = args.num_devices
	self.prefetcher = prefetcher
	self.use_galaxy_no_prefetch = args.is_galaxy and prefetcher is None

	vocab_shards = args.cluster_shape[0] if self.use_galaxy_no_prefetch else self.num_devices
	size_per_device = self.padded_vocab_size // vocab_shards

	tile_size = 32
	max_columns_per_device_ring_mm = math.ceil((max_columns_per_device) / tile_size) * tile_size
	max_columns_per_device_dram_sharded = max_columns_per_device

	self.model_config = args.get_model_config()

	num_splits_ring_mm = math.ceil(size_per_device / max_columns_per_device_ring_mm)
	num_splits_dram_sharded = math.ceil(size_per_device / max_columns_per_device_dram_sharded)

	self.split_sizes_dram_sharded = [min(size_per_device, max_columns_per_device_dram_sharded)] * (
	num_splits_dram_sharded - 1
	)
	self.split_sizes_dram_sharded.append(size_per_device - sum(self.split_sizes_dram_sharded)) # remaining columns

	self.split_sizes_ring_mm = [min(size_per_device, max_columns_per_device_ring_mm)] * (num_splits_ring_mm - 1)
	self.split_sizes_ring_mm.append(size_per_device - sum(self.split_sizes_ring_mm)) # remaining columns

	# Split the output weights
	torch_output_weights = state_dict[f"{state_dict_prefix}output.weight"].permute(1, 0)

	# Pad the output weights to the padded vocab size with zeros
	if self.vocab_size < self.padded_vocab_size:
	padding_size = self.padded_vocab_size - self.vocab_size
	torch_output_weights = torch.cat(
	[
	torch_output_weights,
	torch.zeros(torch_output_weights.shape[0], padding_size, dtype=torch_output_weights.dtype),
	],
	dim=-1,
	)

	self.output_weights_dram_sharded = []
	self.output_weights_ring_mm = []

	self.split_sizes = [self.split_sizes_dram_sharded]
	if self.prefetcher is not None:
	self.split_sizes.append(self.split_sizes_ring_mm)

	for mode, split_sizes in enumerate(self.split_sizes):
	for i, split_size in enumerate(split_sizes):
	# Create a list to store the split tensors for each device
	device_splits = []
	for device in range(vocab_shards):
	start = device * size_per_device + sum(split_sizes[:i])
	end = start + split_size
	device_splits.append(torch_output_weights[:, start:end])

	# Concatenate the splits from all devices
	combined_split = torch.cat(device_splits, dim=-1)

	cache_file_name = _lm_head_cache_file_name(
	weight_cache_path,
	dummy_weights=args.dummy_weights,
	num_splits=len(split_sizes),
	shard=i,
	width=combined_split.shape[-1],
	mode=mode,
	galaxy_2d=self.use_galaxy_no_prefetch and mode == 0,
	mesh_shape=args.cluster_shape,
	)

	def pad_to_power_of_2(n):
	if n <= 0:
	return 1
	return 1 << (n - 1).bit_length()

	if mode == 0:
	if self.use_galaxy_no_prefetch:
	torch_weight = combined_split.unsqueeze(0).unsqueeze(0)
	mesh_mapper = ttnn.ShardTensor2dMesh(mesh_device, dims=(3, 2), mesh_shape=args.cluster_shape)
	memory_config = ttnn.DRAM_MEMORY_CONFIG
	else:
	torch_weight = combined_split
	mesh_mapper = ttnn.ShardTensorToMesh(mesh_device, dim=-1)
	memory_config = args.create_dram_sharded_mem_config(
	k=args.dim, n=math.ceil(combined_split.shape[-1] / self.num_devices)
	)
	self.output_weights_dram_sharded.append(
	ttnn.as_tensor(
	torch_weight,
	device=mesh_device,
	mesh_mapper=mesh_mapper,
	layout=ttnn.TILE_LAYOUT,
	dtype=dtype,
	memory_config=memory_config,
	cache_file_name=cache_file_name,
	)
	)
	else:
	memory_config = args.create_dram_sharded_mem_config(
	k=args.dim,
	n=pad_to_power_of_2(math.ceil(combined_split.shape[-1] / self.num_devices)),
	dram_grid=self.prefetcher.to_core_range_set(self.prefetcher.dram_banks()),
	)
	self.output_weights_ring_mm.append(
	ttnn.as_tensor(
	combined_split,
	device=mesh_device,
	mesh_mapper=ttnn.ShardTensorToMesh(mesh_device, dim=-1),
	layout=ttnn.TILE_LAYOUT,
	dtype=dtype,
	memory_config=memory_config,
	cache_file_name=cache_file_name,
	)
	)

	self.compute_kernel_config = ttnn.WormholeComputeKernelConfig(
	math_fidelity=ttnn.MathFidelity.HiFi2,
	math_approx_mode=False,
	fp32_dest_acc_en=False,
	packer_l1_acc=True,
	)

	def _update_output_weights_dram_sharded(self, weight: ttnn.Tensor) -> None:
	"""In-place replace every chunk of ``self.output_weights_dram_sharded``
	on device (``weight`` is HF ``(1, 1, V, H)``, replicated, TILE, bf16,
	DRAM-interleaved). Mirrors the constructor's host-side pad+transpose+split
	but on device: optional ``ttnn.pad`` to ``padded_vocab_size``,
	``ttnn.transpose``, contiguous per-chunk ``ttnn.slice``, then
	``copy_to_buffer`` reshards into the DRAM-sharded dest (preserving
	addresses). Single-device only; multi-device per-device gather not
	implemented.

	Tile-alignment caveat: ``ttnn.slice`` on TILE requires offsets aligned
	to TILE_SIZE (32); the default Llama configs satisfy this, but assert it
	rather than silently miscompile slices.
	"""
	assert self.num_devices == 1, (
	"LMHead.update for num_devices > 1 is not yet implemented; "
	"the multi-device path needs a per-device interleaved gather."
	)
	tile_size = 32
	assert all(s % tile_size == 0 for s in self.split_sizes_dram_sharded), (
	f"LMHead.update requires every split in self.split_sizes_dram_sharded "
	f"to be a multiple of TILE_SIZE={tile_size}, got "
	f"{self.split_sizes_dram_sharded}; on-device ttnn.slice cannot produce "
	"sub-tile-aligned chunks on TILE_LAYOUT."
	)

	padded = weight
	if self.vocab_size < self.padded_vocab_size:
	pad_amount = self.padded_vocab_size - self.vocab_size
	padded = ttnn.pad(padded, [(0, 0), (0, 0), (0, pad_amount), (0, 0)], value=0.0)

	permuted = ttnn.transpose(padded, -2, -1)

	start = 0
	for i, split_size in enumerate(self.split_sizes_dram_sharded):
	end = start + split_size
	chunk = ttnn.slice(permuted, (0, 0, 0, start), (1, 1, self.args.dim, end))
	copy_to_buffer(chunk, self.output_weights_dram_sharded[i], self.dtype)
	start = end

	def _update_output_weights_ring_mm(self, weight: ttnn.Tensor) -> None:
	"""In-place replace every chunk of ``self.output_weights_ring_mm``. Not
	yet implemented: the ring-mm path (distinct per-chunk split, power-of-two
	column padding, prefetcher-derived DRAM grid) needs its own builder.
	"""
	raise NotImplementedError("LMHead.update for output_weights_ring_mm (prefetcher path) is not yet implemented")

	def update(self, *, weight: ttnn.Tensor) -> None:
	"""In-place replace the on-device LM-head weights via ``ttnn.copy``.

	HF-format input (see ``LLAMA_WEIGHT_TRANSFER.md``): ``weight`` is
	``lm_head.weight``, shape ``(1, 1, vocab_size, hidden_size)``, bf16, TILE,
	DRAM-interleaved, replicated.

	Updates ``self.output_weights_dram_sharded`` on device (no host
	roundtrip). When the prefetcher is enabled the ring-mm mirror must also
	be synced; that path is not implemented. Every chunk keeps its device
	allocation, so captured traces and the prefetcher's recorded addresses
	stay valid.
	"""
	self._update_output_weights_dram_sharded(weight)
	if len(self.output_weights_ring_mm) > 0:
	self._update_output_weights_ring_mm(weight)

	def forward(self, x: ttnn.Tensor, debug_input_torch=None, debug_weight_torch=None):
	outputs = []
	use_prefetcher = self.prefetcher is not None and self.prefetcher.mode == Mode.DECODE
	split_sizes = self.split_sizes_ring_mm if use_prefetcher else self.split_sizes_dram_sharded
	program_configs = [
	self.args.get_lm_head_program_config(split_size, self.prefetcher if use_prefetcher else None)
	for split_size in split_sizes
	]

	output_weights = self.output_weights_ring_mm if use_prefetcher else self.output_weights_dram_sharded

	if self.use_galaxy_no_prefetch:
	x = ttnn.to_memory_config(x, ttnn.DRAM_MEMORY_CONFIG)

	self.lm_head_output_memory_config = self.args.get_lm_head_output_mem_config(
	Mode.DECODE if use_prefetcher else Mode.PREFILL, self.prefetcher if use_prefetcher else None
	)

	for i, (weight, pc) in enumerate(zip(output_weights, program_configs)):
	output = ttnn.linear(
	x,
	weight,
	compute_kernel_config=self.compute_kernel_config,
	program_config=None if self.use_galaxy_no_prefetch else pc,
	memory_config=(
	ttnn.DRAM_MEMORY_CONFIG if self.use_galaxy_no_prefetch else self.lm_head_output_memory_config
	),
	dtype=self.args.lm_head_dtype if hasattr(self.args, "lm_head_dtype") else ttnn.bfloat8_b,
	sub_device_id=self.prefetcher.worker_sub_device_id if use_prefetcher else None,
	)
	output = ttnn.to_memory_config(
	output,
	memory_config=self.args.get_lm_head_sharded_output_mem_config(
	self.prefetcher if use_prefetcher else None
	),
	)

	outputs.append(output)

	ttnn.deallocate(x)

	# Concatenate the outputs
	# outputs shape: a list of tensors, each tensor is 1,1,32,size_per_device per device
	# Decode traces retain these logits across requests. Honor the model's placement
	# so large vocabularies do not occupy L1 needed by the next eager prefill.
	output = ttnn.concat(
	outputs,
	dim=-1,
	memory_config=(
	self.args.model_config.get("LM_HEAD_OUTPUT_MEMCFG", ttnn.L1_MEMORY_CONFIG)
	if not use_prefetcher
	else ttnn.DRAM_MEMORY_CONFIG
	),
	sub_core_grids=self.prefetcher.all_worker_cores_range_set if use_prefetcher else None,
	)

	# Only use reshard mem config for ring_mm mode
	if use_prefetcher:
	output = ttnn.to_memory_config(
	output,
	memory_config=self.args.get_lm_head_reshard_mem_config(self.prefetcher),
	)

	output = tt_all_reduce(
	output,
	self.mesh_device,
	self.tt_ccl,
	cluster_axis=1,
	dim=3 if self.args.is_galaxy else 0,
	memory_config=output.memory_config(),
	dtype=self.args.ccl_dtype,
	sharded=False,
	use_composite=True,
	subdevice_id=self.prefetcher.worker_sub_device_id if use_prefetcher else None,
	)

	return output