from megatron.core import parallel_state

# 분산 환경 초기화
parallel_state.initialize_model_parallel(
    tensor_model_parallel_size=8,
    pipeline_model_parallel_size=1,
    virtual_pipeline_model_parallel_size=None,
)

# TP 그룹 조회
tp_group = parallel_state.get_tensor_model_parallel_group()
tp_rank = parallel_state.get_tensor_model_parallel_rank()
tp_world_size = parallel_state.get_tensor_model_parallel_world_size()

# megatron/core/tensor_parallel/mappings.py:197-214
class _CopyToModelParallelRegion(torch.autograd.Function):
    @staticmethod
    def forward(ctx, input_, group):
        """Forward: identity (통신 없음)"""
        ctx.group = group
        return input_

    @staticmethod
    def backward(ctx, grad_output):
        """Backward: all-reduce gradients"""
        return _reduce(grad_output, ctx.group), None

# megatron/core/tensor_parallel/mappings.py:217-233
class _ReduceFromModelParallelRegion(torch.autograd.Function):
    @staticmethod
    def forward(ctx, input_, group):
        """Forward: all-reduce across TP ranks"""
        return _reduce(input_, group)

    @staticmethod
    def backward(ctx, grad_output):
        """Backward: identity (통신 없음)"""
        return grad_output, None

# megatron/core/tensor_parallel/mappings.py:236-273
class _ScatterToModelParallelRegion(torch.autograd.Function):
    @staticmethod
    def forward(ctx, input_, group):
        """Forward: split along last dim"""
        ctx.group = group
        return _split_along_last_dim(input_, group)

    @staticmethod
    def backward(ctx, grad_output):
        """Backward: gather along last dim"""
        return _gather_along_last_dim(grad_output, ctx.group), None

class _GatherFromModelParallelRegion(torch.autograd.Function):
    @staticmethod
    def forward(ctx, input_, group):
        """Forward: gather along last dim"""
        ctx.group = group
        return _gather_along_last_dim(input_, group)

    @staticmethod
    def backward(ctx, grad_output):
        """Backward: split along last dim"""
        return _split_along_last_dim(grad_output, ctx.group), None

# megatron/core/tensor_parallel/layers.py:494-627
class LinearWithGradAccumulationAndAsyncCommunication(torch.autograd.Function):
    @staticmethod
    def backward(ctx, grad_output):
        input, weight = ctx.saved_tensors
        tp_group = ctx.tp_group

        # 1. 입력 그래디언트 계산
        grad_input = grad_output.matmul(weight)

        # 2. 비동기 all-reduce 시작 (통신)
        if ctx.allreduce_dgrad:
            handle = torch.distributed.all_reduce(
                grad_input, group=tp_group, async_op=True
            )
            # Here we rely on CUDA_DEVICE_MAX_CONNECTIONS=1 to ensure that the
            # all-reduce is scheduled before the weight gradient computation

        # 3. 가중치 그래디언트 계산 (계산) - 통신과 동시에!
        if ctx.gradient_accumulation_fusion:
            # CUDA 커널로 main_grad에 직접 누적 → 중간 텐서 할당 제거
            if weight.main_grad.dtype == torch.float32:
                fused_weight_gradient_mlp_cuda.wgrad_gemm_accum_fp32(
                    total_input, grad_output, weight.main_grad
                )
            # ...
        else:
            grad_weight = grad_output.t().matmul(total_input)

        # 4. all-reduce 완료 대기
        if ctx.allreduce_dgrad:
            handle.wait()

        return grad_input, grad_weight, ...

# megatron/core/tensor_parallel/utils.py:97-121
class VocabUtility:
    @staticmethod
    def vocab_range_from_global_vocab_size(
        global_vocab_size: int, rank: int, world_size: int
    ) -> Sequence[int]:
        """Vocab range from global vocab size."""
        per_partition_vocab_size = divide(global_vocab_size, world_size)
        return VocabUtility.vocab_range_from_per_partition_vocab_size(
            per_partition_vocab_size, rank, world_size
        )

    @staticmethod
    def vocab_range_from_per_partition_vocab_size(
        per_partition_vocab_size: int, rank, world_size: int
    ) -> Sequence[int]:
        index_f = rank * per_partition_vocab_size
        index_l = index_f + per_partition_vocab_size
        return index_f, index_l

# megatron/core/tensor_parallel/cross_entropy.py
class _VocabParallelCrossEntropy(torch.autograd.Function):
    @staticmethod
    def forward(ctx, vocab_parallel_logits, target):
        # 1. logits_max: 수치 안정성을 위한 all-reduce (MAX)
        logits_max = torch.max(vocab_parallel_logits, dim=-1)[0]
        torch.distributed.all_reduce(logits_max, op=ReduceOp.MAX)

        # 2. predicted_logits: 타겟 토큰의 logit만 all-reduce (SUM)
        # (각 랭크는 자신의 vocab 파티션에 해당하는 타겟만 처리)

        # 3. sum_exp_logits: 분할 함수 계산을 위한 all-reduce (SUM)

        # 최종: loss = log(sum_exp) - predicted_logits

from nanotron.config import ParallelismArgs
from nanotron.parallel.context import ParallelContext

# 병렬화 설정
parallelism = ParallelismArgs(
    dp=1,
    pp=1,
    tp=8,
    tp_mode="ALL_REDUCE",  # 또는 "REDUCE_SCATTER"
    tp_linear_async_communication=True,
)

# ParallelContext 초기화
parallel_context = ParallelContext(
    tensor_parallel_size=parallelism.tp,
    pipeline_parallel_size=parallelism.pp,
    data_parallel_size=parallelism.dp,
)

# src/nanotron/parallel/tensor_parallel/enum.py
class TensorParallelLinearMode(Enum):
    ALL_REDUCE = "all_reduce"
    REDUCE_SCATTER = "reduce_scatter"

if tp_mode is TensorParallelLinearMode.ALL_REDUCE:
    gathered_tensor = tensor  # 통신 없음, 입력이 이미 복제됨
    return F.linear(gathered_tensor, weight, bias)

elif tp_mode is TensorParallelLinearMode.REDUCE_SCATTER:
    # Forward: AllGather
    handle = dist.all_gather_into_tensor(gathered_tensor, tensor, group=group, async_op=True)
    # Backward: ReduceScatter

# src/nanotron/parallel/tensor_parallel/functional.py:305-327
# AllGather 비동기 시작
handle = dist.all_gather_into_tensor(gathered_tensor, tensor, group=group, async_op=True)

# 출력 텐서를 before/same/after로 분할
before_shard, same_device_shard, after_shard = torch.split(
    gathered_output,
    split_size_or_sections=[
        sharded_batch_size * current_rank,
        sharded_batch_size,  # 자신의 데이터
        sharded_batch_size * (group_size - current_rank - 1),
    ],
    dim=0,
)

# AllGather 완료 전에 자신의 shard 먼저 계산
torch.mm(
    input=tensor.view(first_dims, hidden_size),
    mat2=weight.t(),
    out=same_device_shard.view(first_dims, output_size),
)

# AllGather 완료 대기
handle.wait()

# 나머지 shard 계산
if before_shard.numel() > 0:
    torch.mm(
        input=gathered_tensor[: sharded_batch_size * current_rank].view(first_dims, hidden_size),
        mat2=weight.t(),
        out=before_shard.view(first_dims, output_size),
    )
# after_shard도 동일하게 처리

# src/nanotron/parallel/tensor_parallel/nn.py
class TensorParallelColumnLinear(nn.Linear):
    def __init__(self, ..., tp_recompute_allgather=False):
        self.tp_recompute_allgather = tp_recompute_allgather

if tp_recompute_allgather:
    gathered_tensor = MemoryBuffer().get("allgather", ...)  # 버퍼 재사용
    ctx.save_for_backward(tensor, weight)  # sharded tensor만 저장

if ctx.tp_recompute_allgather:
    unsharded_tensor = MemoryBuffer().get("allgather", ...)
    handle = dist.all_gather_into_tensor(unsharded_tensor, tensor, group=group, async_op=True)
    # AllGather 다시 수행

qkv_contiguous_chunks = (
    config.num_attention_heads * self.d_qk,      # Q 청크
    config.num_key_value_heads * self.d_qk,      # K 청크
    config.num_key_value_heads * self.d_qk,      # V 청크
)

self.qkv_proj = TensorParallelColumnLinear(
    hidden_size, q_out + 2*kv_out,
    contiguous_chunks=qkv_contiguous_chunks,
)

# src/nanotron/parallel/context.py
# [expert_parallel, pipeline_parallel, data_parallel, context_parallel, tensor_parallel]

self.tp_pg         # TP group
self.dp_pg         # DP group
self.pp_pg         # PP group
self.cp_pg         # Context Parallel group (시퀀스 병렬화)
self.ep_pg         # Expert Parallel group (MoE용)

import deepspeed
from transformers import AutoModelForCausalLM

# 모델 로드
model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-2-7b-hf")

# DeepSpeed 초기화 (AutoTP 자동 적용)
model_engine, optimizer, _, _ = deepspeed.initialize(
    model=model,
    config="ds_config.json",
)

# 또는 수동으로 TP 적용
from deepspeed.module_inject import auto_tp
model = auto_tp.AutoTP(model, tp_size=8)

# deepspeed/module_inject/auto_tp.py:307-338
def tp_parser(model):
    for module in module_list:
        for key, submodule in module._modules.items():
            if isinstance(submodule, nn.Linear):
                layer_list = layer_list + ["." + key]
            elif isinstance(submodule, nn.LayerNorm) or key in norm_layer_name_list:
                layer_list = layer_list + ["ln"]

        for i, layer in enumerate(layer_list):
            if layer == 'ln':
                if layer_list[i - 1] != 'ln':
                    gem_list.append(layer_list[i - 1])  # LN 직전 = all-reduce 필요
            elif 'out_proj' in layer:
                gem_list.append(layer)
            elif 'down_proj' in layer:
                gem_list.append(layer)

# deepspeed/module_inject/fusedqkv_utils.py:34-46
fused_type_dict = {
    'CodeGenBlock': 'codegentype',   # [q(1),q(2),...,k(1),k(2),...,v(1),v(2),...]
    'BloomBlock': 'bloomtype',        # [q(1),k(1),v(1),q(2),k(2),v(2),...]
    'GLMBlock': 'glmtype',            # [Q,Q,...,K,K,...,V,V,...]
    "MPTBlock": 'glmtype',
    "BaichuanLayer": 'glmtype',
    "QWenBlock": 'qwentype',
    "FalconDecoderLayer": 'bloomtype',
    "GPTBigCodeBlock": 'bigcodetype',
    "Phi3DecoderLayer": "phi3type",   # Rotary 임베딩 위치 분리
}

# deepspeed/module_inject/layers.py
class SubParamLinearLayer(TensorParallel_Layer):
    def __init__(self, module, mp_group, shape, partition_dim=0):
        # shape = ((q_size, k_size, v_size), -1) for GQA
        # 각 서브파라미터를 독립적으로 분할

shape = ((4096, 1024, 1024), -1)  # Q: 4096, K: 1024, V: 1024
# 각 sub-param을 독립적으로 분할
sub_params = torch.split(tensor, subparam_sizes, dim=partition_dim)
partitioned = [torch.chunk(sp, tp_size, dim=0)[tp_idx] for sp in sub_params]

# deepspeed/module_inject/tp_shard.py:47-67
def get_shard_size(total_size, mp_size, name=None, rank=None):
    if total_size >= tp_grain_size:
        grain_size = total_size // tp_grain_size
        return (grain_size // mp_size + (1 if rank < (grain_size % mp_size) else 0)) * tp_grain_size
    else:
        return total_size // mp_size + (1 if rank < (total_size % mp_size) else 0)

from torchtitan.distributed import ParallelDims
from torch.distributed.device_mesh import init_device_mesh

# 병렬화 차원 정의
parallel_dims = ParallelDims(
    tp=8,
    pp=1,
    dp_shard=-1,  # auto
    dp_replicate=1,
    cp=1,
)

# DeviceMesh 초기화
world_mesh = parallel_dims.build_mesh(device_type="cuda")

# TP용 sub-mesh 추출
tp_mesh = world_mesh["tp"]

# torchtitan/models/llama3/infra/parallelize.py:161-248
from torch.distributed.tensor import Replicate, Shard
from torch.distributed.tensor.parallel import (
    ColwiseParallel, RowwiseParallel, SequenceParallel,
    parallelize_module, PrepareModuleInput,
)

def apply_tp(model, tp_mesh, loss_parallel, enable_float8_tensorwise_tp, cp_enabled):
    # 1. 임베딩, 정규화, 출력 레이어 병렬화
    parallelize_module(
        model,
        tp_mesh,
        {
            "tok_embeddings": RowwiseParallel(
                input_layouts=Replicate(),
                output_layouts=Shard(1),
            ),
            "norm": SequenceParallel(),
            "output": ColwiseParallel(
                input_layouts=Shard(1),
                output_layouts=Shard(-1) if loss_parallel else Replicate(),
                use_local_output=not loss_parallel,
            ),
        },
    )

    # 2. 각 Transformer 블록 병렬화
    for transformer_block in model.layers.values():
        layer_plan = {
            "attention_norm": SequenceParallel(),
            "attention": PrepareModuleInput(
                input_layouts=(Shard(1), None, None, None),      # 현재: sequence dim 분할
                desired_input_layouts=(Replicate(), None, None, None),  # 목표: 복제
            ),
            "attention.wq": ColwiseParallel(),
            "attention.wk": ColwiseParallel(),
            "attention.wv": ColwiseParallel(),
            "attention.wo": RowwiseParallel(output_layouts=Shard(1)),
            "ffn_norm": SequenceParallel(),
            "feed_forward": PrepareModuleInput(
                input_layouts=(Shard(1),),
                desired_input_layouts=(Replicate(),),
            ),
            "feed_forward.w1": ColwiseParallel(),
            "feed_forward.w2": RowwiseParallel(output_layouts=Shard(1)),
            "feed_forward.w3": ColwiseParallel(),
        }
        parallelize_module(transformer_block, tp_mesh, layer_plan)

# torchtitan/distributed/utils.py
if enable_loss_parallel:
    stack.enter_context(torch.distributed.tensor.parallel.loss_parallel())

# torchtitan/models/llama3/infra/parallelize.py:192-210
if enable_float8_tensorwise_tp:
    from torchao.float8.float8_tensor_parallel import (
        Float8ColwiseParallel,
        Float8RowwiseParallel,
        PrepareFloat8ModuleInput,
    )
    rowwise_parallel, colwise_parallel, prepare_module_input = (
        Float8RowwiseParallel,
        Float8ColwiseParallel,
        PrepareFloat8ModuleInput,
    )

from torch.distributed.tensor import Replicate, Shard

Replicate()   # 전체 텐서를 모든 랭크에 복제
Shard(0)      # 0번 차원으로 분할 (batch)
Shard(1)      # 1번 차원으로 분할 (sequence/hidden)
Shard(-1)     # 마지막 차원으로 분할 (vocab dim)

# Async TP 활성화
torch._inductor.config._micro_pipeline_tp = True

from vllm import LLM, SamplingParams

# 모델 로드 (TP 자동 적용)
llm = LLM(
    model="meta-llama/Llama-3.1-8B-Instruct",
    tensor_parallel_size=8,
    max_model_len=8192,
    gpu_memory_utilization=0.9,
)

# 추론
sampling_params = SamplingParams(temperature=0.7, max_tokens=256)
outputs = llm.generate(["Hello, world!"], sampling_params)

# vllm/distributed/parallel_state.py:276-505
class GroupCoordinator:
    """PyTorch ProcessGroup wrapper for a group of processes."""

    def __init__(self, group_ranks, local_rank, torch_distributed_backend, ...):
        self.unique_name = _get_unique_name(group_name)
        _register_group(self)  # custom op에서 조회할 수 있도록 등록

        # CPU와 device 통신 그룹 분리
        self.cpu_group = torch.distributed.new_group(ranks, backend="gloo")
        self.device_group = torch.distributed.new_group(ranks, backend=torch_distributed_backend)

    def all_reduce(self, input_: torch.Tensor) -> torch.Tensor:
        # Bypass the function if we are using only 1 GPU.
        if self.world_size == 1:
            return input_

        if self.use_custom_op_call:
            return torch.ops.vllm.all_reduce(input_, group_name=self.unique_name)
        else:
            return self._all_reduce_out_place(input_)

# vllm/distributed/parallel_state.py:248-273
direct_register_custom_op(
    op_name="all_reduce",
    op_func=all_reduce,
    fake_impl=all_reduce_fake,  # torch.compile용 shape 추론
)

direct_register_custom_op(
    op_name="reduce_scatter",
    op_func=reduce_scatter,
    fake_impl=reduce_scatter_fake,
)

direct_register_custom_op(
    op_name="all_gather",
    op_func=all_gather,
    fake_impl=all_gather_fake,
)

# vllm/model_executor/layers/linear.py:551-586
class ColumnParallelLinear(LinearBase):
    def weight_loader(self, param, loaded_weight):
        output_dim = getattr(param, "output_dim", None)
        is_sharded_weight = getattr(param, "is_sharded_weight", False)

        if output_dim is not None and not is_sharded_weight:
            shard_size = param_data.shape[output_dim]
            start_idx = self.tp_rank * shard_size
            # 로드 시점에 weight를 분할
            loaded_weight = loaded_weight.narrow(output_dim, start_idx, shard_size)

        param_data.copy_(loaded_weight)

# vllm/model_executor/layers/linear.py:954-962
class QKVParallelLinear(ColumnParallelLinear):
    def __init__(self, hidden_size, head_size, total_num_heads, total_num_kv_heads, ...):
        tp_size = get_tensor_model_parallel_world_size()

        if tp_size >= self.total_num_kv_heads:
            # KV 헤드 < TP 크기: KV 헤드 복제
            self.num_kv_heads = 1
            self.num_kv_head_replicas = divide(tp_size, self.total_num_kv_heads)
        else:
            # KV 헤드 >= TP 크기: KV 헤드 분할
            self.num_kv_heads = divide(self.total_num_kv_heads, tp_size)
            self.num_kv_head_replicas = 1

# vllm/model_executor/layers/linear.py (RowParallelLinear)
def forward(self, input_):
    if self.input_is_parallel:
        input_parallel = input_  # scatter 스킵, 이미 분할되어 있음
    else:
        splitted = split_tensor_along_last_dim(input_, self.tp_size)
        input_parallel = splitted[self.tp_rank]

    output_parallel = self.quant_method.apply(self, input_parallel, bias)

    if self.reduce_results and self.tp_size > 1:
        output = tensor_model_parallel_all_reduce(output_parallel)

import sglang as sgl

# 엔진 초기화
engine = sgl.Engine(
    model_path="meta-llama/Llama-3.1-8B-Instruct",
    tp_size=8,
)

# 추론
outputs = engine.generate(
    prompt="Hello, world!",
    sampling_params={"temperature": 0.7, "max_new_tokens": 256},
)

engine = sgl.Engine(
    model_path="meta-llama/Llama-3.1-8B-Instruct",
    tp_size=8,
    dp_size=2,                      # DP+TP 이중 병렬화
    enable_flashinfer_allreduce_fusion=True,  # Hopper/Blackwell 최적화
)

# sglang/srt/layers/communicator.py:102-120
class ScatterMode(Enum):
    """
    Suppose we have TP=4, DP=2, enable-dp-attention, and the system handles seq a,b,c,d
    Model input/output: [ab, ab, cd, cd] for four ranks respectively
    """
    SCATTERED = auto()      # [a, b, c, d] - 각 랭크가 자신만
    TP_ATTN_FULL = auto()   # [ab, ab, cd, cd] - TP attn 그룹 내 복제
    FULL = auto()           # [abcd, abcd, abcd, abcd] - 전체 복제

    @staticmethod
    def model_input_output():
        """The scatter mode for model forward pass input and output data"""
        if is_nsa_enable_prefill_cp():
            return ScatterMode.SCATTERED
        return ScatterMode.TP_ATTN_FULL

# sglang/srt/layers/communicator.py:89-99
FUSE_ALLREDUCE_MAX_BATCH_SIZE = 2048

def apply_flashinfer_allreduce_fusion(batch_size: int):
    return (
        (_is_sm90_supported or _is_sm100_supported)  # Hopper/Blackwell
        and _is_flashinfer_available
        and batch_size > 0
        and batch_size <= FUSE_ALLREDUCE_MAX_BATCH_SIZE
        and not is_dp_attention_enabled()
        and get_global_server_args().enable_flashinfer_allreduce_fusion
    )

# sglang/srt/layers/communicator.py:336-380
class LayerCommunicator:
    def __init__(
        self,
        layer_scatter_modes: LayerScatterModes,
        input_layernorm: torch.nn.Module,
        post_attention_layernorm: torch.nn.Module,
        allow_reduce_scatter: bool = False,
        is_last_layer: bool = False,
        qkv_latent_func: Optional[Callable] = None,
    ):
        self.layer_scatter_modes = layer_scatter_modes
        # 레이어별로 최적의 통신 패턴 결정
        self._communicate_simple_fn = CommunicateSimpleFn.get_fn(
            input_mode=self.layer_scatter_modes.layer_input_mode,
            output_mode=self.layer_scatter_modes.attn_mode,
            context=self._context,
        )

@dataclass
class LayerScatterModes:
    layer_input_mode: ScatterMode    # 레이어 입력
    attn_mode: ScatterMode           # Attention 계산
    mlp_mode: ScatterMode            # MLP 계산
    middle_residual_mode: ScatterMode # 중간 residual
    layer_output_mode: ScatterMode   # 레이어 출력

# sglang/srt/layers/dp_attention.py
def get_attention_tp_group() -> GroupCoordinator:
    return _ATTN_TP_GROUP  # Attention 전용 TP 그룹

def get_attention_dp_size() -> int:
    return _ATTN_DP_SIZE  # DP 크기

# sglang/srt/layers/communicator.py:608-641
@dataclass
class CommunicateContext:
    process_group_sizes: Dict[ScatterMode, int]
    attn_tp_rank: int
    attn_tp_size: int
    attn_dp_size: int
    tp_size: int
    tp_rank: int

    @classmethod
    def init_new(cls):
        process_group_sizes = {
            ScatterMode.SCATTERED: 1,
            ScatterMode.TP_ATTN_FULL: attn_tp_size,
            ScatterMode.FULL: tp_size,
        }
        return cls(...)

def vocab_parallel_cross_entropy(vocab_parallel_logits, target):
    # 각 랭크: [batch*seq, vocab/p] logits 보유

    # 1. 수치 안정성: max logit을 all-reduce (MAX)
    logits_max = vocab_parallel_logits.max(dim=-1)
    all_reduce(logits_max, op=MAX)

    # 2. 타겟 logit: 해당 랭크만 값을 가짐, all-reduce (SUM)
    # 3. 분할 함수: exp(logits).sum()을 all-reduce (SUM)

    # 통신량: O(batch * seq) - vocab 차원 제거됨