Automated Parallelization System and Infrastructure for Multiple Ecosystems
from dataclasses import dataclass


@dataclass
class GPTCase:
    batch_size: int = 4
    seq_size: int = 1024
    num_layers: int = 1
    hidden_dim: int = 12288
    num_heads: int = 48
    dropout_rate: float = 0.0
    use_bias: bool = True
    dtype = "float32"


@dataclass
class ResNetCase:
    batch_size: int = 128


@dataclass
class GATCase:
    num_node: int = 4096
    in_feature: int = 12288
    out_feature: int = 12288

benchmark/bench_case.py
alibaba/easydist
# mpirun --mca btl tcp,self -np 2 python ./benchmark/bench_jax.py

import os
import sys
import logging

import jax
from jax.experimental import mesh_utils
from jax.sharding import Mesh
from mpi4py import MPI

from easydist.jax import get_opt_strategy, easydist_shard, set_device_mesh
from easydist.jax.api import shard_module, to_shape_array
from easydist import easydist_setup
from easydist.utils.timer import EDTimer

sys.path.append(os.path.dirname(os.path.abspath(__file__)))
from benchmark.bench_case import GPTCase, ResNetCase, GATCase
from benchmark.jax.model import GATLayer
from benchmark.jax.model.gpt import GPTSimple
from benchmark.jax.model.wresnet import wresnet50

logging.basicConfig(format='%(asctime)s - %(levelname)s - %(message)s',
                    datefmt='%m/%d %H:%M:%S',
                    level=logging.INFO)


def get_gpt_case():
    case = GPTCase()
    model = GPTSimple(case)

    root_key = jax.random.PRNGKey(seed=0)
    main_key, params_key = jax.random.split(key=root_key)
    input_ = jax.random.normal(
        main_key, (case.batch_size, case.seq_size, case.hidden_dim))  # Dummy input data
    variables = model.init(params_key, input_, deterministic=True)
    params = variables['params']

    def train_step(params, input_):
        lr = 0.0001

        def loss_fn(params):
            dropout_key = jax.random.PRNGKey(seed=0)
            return model.apply({
                'params': params
            },
                               input_,
                               deterministic=False,
                               rngs={
                                   'dropout': dropout_key
                               }).mean()

        grad_fn = jax.grad(loss_fn)
        grads = grad_fn(params)
        params = jax.tree_map(lambda x, y: x - lr * y, params, grads)
        return params

    return train_step, [params, input_]


def get_resnet_case():
    case = ResNetCase()
    model = wresnet50()

    key1, key2 = jax.random.split(jax.random.PRNGKey(0), num=2)
    input_ = jax.random.normal(key1, (case.batch_size, 224, 224, 3))  # Dummy input data
    variables = model.init(key2, input_)  # Initialization call
    params, batch_stats = variables['params'], variables['batch_stats']

    def train_step(params, batch_stats, input_):
        lr = 0.0001

        def loss_fn(params, batch_stats):
            out_, batch_stats = model.apply({
                'params': params,
                'batch_stats': batch_stats
            },
                                            input_,
                                            mutable=['batch_stats'])
            return out_.mean()

        grad_fn = jax.grad(loss_fn)
        grads = grad_fn(params, batch_stats)
        params = jax.tree_map(lambda x, y: x - lr * y, params, grads)
        return params

    return train_step, [params, batch_stats, input_]


def get_gat_case():

    case = GATCase()
    model = GATLayer(case.in_feature, case.out_feature)

    key1, key2, key3 = jax.random.split(jax.random.PRNGKey(0), num=3)
    h = jax.random.normal(key1, (case.num_node, case.in_feature))  # Dummy input data
    adj = jax.random.normal(key2, (case.num_node, case.num_node))  # Dummy input data
    variables = model.init(key3, h, adj)  # Initialization call
    params = variables['params']

    def train_step(params, h, adj):
        lr = 0.0001

        def loss_fn(params):
            return model.apply({'params': params}, h, adj).mean()

        grad_fn = jax.grad(loss_fn)
        grads = grad_fn(params)
        params = jax.tree_map(lambda x, y: x - lr * y, params, grads)
        return params

    return train_step, [params, h, adj]


def bench_naive(func, args):

    jit_func = jax.jit(func)

    def train_step():
        jit_func(*args)

    timer = EDTimer(train_step, in_ms=False)

    elaps_time = timer.time()

    print(f"Time: {elaps_time}")


def bench_easydist(func, args):
    size = jax.device_count()
    devices = mesh_utils.create_device_mesh((1, size))
    mesh = Mesh(devices, axis_names=('a', 'b'))

    set_device_mesh(mesh)

    opt_strategy = get_opt_strategy(func, *args)

    shard_func = jax.jit(easydist_shard(func, opt_strategy))

    flatten_args, specs = jax.tree_util.tree_flatten(args)

    flatten_args = shard_module(flatten_args)

    shard_args = jax.tree_util.tree_unflatten(specs, flatten_args)

    def train_step():
        with jax.spmd_mode('allow_all'):
            shard_func(*shard_args)

    # warmup
    train_step()

    timer = EDTimer(train_step, in_ms=False)

    elaps_time = timer.time()

    print(f"Time: {elaps_time}")


def main():
    # setup easydist
    easydist_setup(backend="jax", device="cuda")

    func, args = get_gat_case()
    args = jax.tree_map(to_shape_array, args)

    bench_easydist(func, args)


if __name__ == '__main__':
    main()

benchmark/jax/bench_jax.py
alibaba/easydist
# python ./benchmark/bench_jax_alpa.py

import logging
import os
import sys

import alpa
import jax

os.environ["EASYDIST_DEVICE"] = "cuda"
os.environ["EASYDIST_BACKEND"] = "jax"

from easydist.utils.timer import EDTimer

sys.path.append(os.path.dirname(os.path.abspath(__file__)))
from benchmark.jax.model.gpt import GPTSimple
from benchmark.jax.model.wresnet import resnet18
from benchmark.jax.model.gat import GATLayer
from benchmark.bench_case import GPTCase, ResNetCase, GATCase

logging.basicConfig(format='%(asctime)s - %(levelname)s - %(message)s',
                    datefmt='%m/%d %H:%M:%S',
                    level=logging.DEBUG)


def get_gpt_case():
    case = GPTCase()
    model = GPTSimple(case)

    root_key = jax.random.PRNGKey(seed=0)
    main_key, params_key = jax.random.split(key=root_key)
    input_ = jax.random.normal(
        main_key, (case.batch_size, case.seq_size, case.hidden_dim))  # Dummy input data
    variables = model.init(params_key, input_, deterministic=True)
    params = variables['params']

    # DataParallel() Zero3Parallel() Zero2Parallel()
    @alpa.parallelize(method=alpa.ShardParallel())
    def train_step(params, input_):
        lr = 0.0001

        def loss_fn(params):
            dropout_key = jax.random.PRNGKey(seed=0)
            return model.apply({
                'params': params
            },
                               input_,
                               deterministic=False,
                               rngs={
                                   'dropout': dropout_key
                               }).mean()

        grad_fn = jax.grad(loss_fn)
        grads = grad_fn(params)
        params = jax.tree_map(lambda x, y: x - lr * y, params, grads)
        return params

    return train_step, [params, input_]


def get_resnet_case():
    case = ResNetCase()
    model = resnet18()

    key1, key2 = jax.random.split(jax.random.PRNGKey(0), num=2)
    input_ = jax.random.normal(key1, (case.batch_size, 224, 224, 3))  # Dummy input data
    variables = model.init(key2, input_)  # Initialization call
    params, batch_stats = variables['params'], variables['batch_stats']

    @alpa.parallelize(method=alpa.ShardParallel())
    def train_step(params, batch_stats, input_):
        lr = 0.0001

        def loss_fn(params, batch_stats):
            out_, batch_stats = model.apply({
                'params': params,
                'batch_stats': batch_stats
            },
                                            input_,
                                            mutable=['batch_stats'])
            return out_.mean()

        grad_fn = jax.grad(loss_fn)
        grads = grad_fn(params, batch_stats)
        params = jax.tree_map(lambda x, y: x - lr * y, params, grads)
        return params

    return train_step, [params, batch_stats, input_]


def get_gat_case():

    case = GATCase()
    model = GATLayer(case.in_feature, case.out_feature)

    key1, key2, key3 = jax.random.split(jax.random.PRNGKey(0), num=3)
    h = jax.random.normal(key1, (case.num_node, case.in_feature))  # Dummy input data
    adj = jax.random.normal(key2, (case.num_node, case.num_node))  # Dummy input data
    variables = model.init(key3, h, adj)  # Initialization call
    params = variables['params']

    @alpa.parallelize(method=alpa.ShardParallel())
    def train_step(params, h, adj):
        lr = 0.0001

        def loss_fn(params):
            return model.apply({'params': params}, h, adj).mean()

        grad_fn = jax.grad(loss_fn)
        grads = grad_fn(params)
        params = jax.tree_map(lambda x, y: x - lr * y, params, grads)
        return params

    return train_step, [params, h, adj]


def bench_alpa(func, args):

    def train_step():
        func(*args)

    timer = EDTimer(train_step, in_ms=False)

    elaps_time = timer.time()

    print(f"Time: {elaps_time}")


def main():
    os.environ["NVIDIA_TF32_OVERRIDE"] = "0"
    print(jax.devices())

    func, args = get_gat_case()

    bench_alpa(func, args)


if __name__ == '__main__':
    main()

benchmark/jax/bench_jax_alpa.py
alibaba/easydist
# python ./benchmark/bench_jax_dp.py

import os
import sys
import logging
from functools import partial

import jax

from easydist import easydist_setup
from easydist.utils.timer import EDTimer

sys.path.append(os.path.dirname(os.path.abspath(__file__)))
from benchmark.bench_case import GPTCase, ResNetCase
from benchmark.jax.model.gpt import GPTSimple
from benchmark.jax.model.wresnet import resnet18


logging.basicConfig(format='%(asctime)s - %(levelname)s - %(message)s',
                    datefmt='%m/%d %H:%M:%S',
                    level=logging.INFO)


def get_gpt_case():
    case = GPTCase()
    model = GPTSimple(case)

    root_key = jax.random.PRNGKey(seed=0)
    main_key, params_key = jax.random.split(key=root_key)
    input_ = jax.random.normal(
        main_key, (case.batch_size, case.seq_size, case.hidden_dim))  # Dummy input data
    variables = model.init(params_key, input_, deterministic=True)
    params = variables['params']

    @partial(jax.pmap, axis_name="batch")
    def train_step(params, input_):
        lr = 0.0001

        def loss_fn(params):
            dropout_key = jax.random.PRNGKey(seed=0)
            return model.apply({
                'params': params
            },
                               input_,
                               deterministic=False,
                               rngs={
                                   'dropout': dropout_key
                               }).mean()

        grad_fn = jax.grad(loss_fn)
        grads = grad_fn(params)
        grads = jax.lax.pmean(grads, axis_name="batch")
        params = jax.tree_map(lambda x, y: x - lr * y, params, grads)
        return params

    devices = jax.local_devices()
    params = jax.device_put_replicated(params, devices)

    def shard_batch(x):
        x = x.reshape((len(devices), -1) + x.shape[1:])
        return jax.device_put_sharded(list(x), devices)

    input_ = jax.tree_map(shard_batch, input_)

    return train_step, [params, input_]


def get_resnet_case():
    case = ResNetCase()
    model = resnet18()

    key1, key2 = jax.random.split(jax.random.PRNGKey(0), num=2)
    input_ = jax.random.normal(key1, (case.batch_size, 224, 224, 3))  # Dummy input data
    variables = model.init(key2, input_)  # Initialization call
    params, batch_stats = variables['params'], variables['batch_stats']

    @partial(jax.pmap, axis_name="batch")
    def train_step(params, batch_stats, input_):
        lr = 0.0001

        def loss_fn(params, batch_stats):
            out_, batch_stats = model.apply({
                'params': params,
                'batch_stats': batch_stats
            },
                                            input_,
                                            mutable=['batch_stats'])
            return out_.mean()

        grad_fn = jax.grad(loss_fn)
        grads = grad_fn(params, batch_stats)
        grads = jax.lax.pmean(grads, axis_name="batch")
        params = jax.tree_map(lambda x, y: x - lr * y, params, grads)
        return params

    devices = jax.local_devices()
    params = jax.device_put_replicated(params, devices)
    batch_stats = jax.device_put_replicated(batch_stats, devices)

    def shard_batch(x):
        x = x.reshape((len(devices), -1) + x.shape[1:])
        return jax.device_put_sharded(list(x), devices)

    input_ = jax.tree_map(shard_batch, input_)

    return train_step, [params, batch_stats, input_]


def bench_pmap_dp(func, args):

    def train_step():
        func(*args)

    timer = EDTimer(train_step, in_ms=False)

    elaps_time = timer.time()

    print(f"Time: {elaps_time}")


def main():
    # setup easydist
    easydist_setup(backend="jax", device="cuda")

    print(jax.devices())

    func, args = get_gpt_case()

    bench_pmap_dp(func, args)


if __name__ == '__main__':
    main()

benchmark/jax/bench_jax_dp.py
alibaba/easydist
from .gpt import GPT, GPTSimple, GPTBlock, GPTConfig
from .resnet import ResNet, ResNet18
from .wresnet import resnet18, resnet34, resnet50, wresnet50, wresnet101
from .gat import GATLayer

__all__ = [
    "GPT", "GPTSimple", "GPTBlock", "GPTConfig", "ResNet", "ResNet18", "resnet18", "resnet34",
    "resnet50", "wresnet50", "wresnet101", "GATLayer"
]

benchmark/jax/model/__init__.py
alibaba/easydist
import jax
from flax import linen as nn


class GATLayer(nn.Module):
    in_features: int
    out_features: int

    @nn.compact
    def __call__(self, h, adj):
        wh = nn.Dense(features=self.out_features)(h)
        wh1 = nn.Dense(features=1)(wh)
        wh2 = nn.Dense(features=1)(wh)
        e = nn.leaky_relu(wh1 + wh2.T)

        zero_vec = -10e10 * jax.numpy.ones_like(e)
        attention = jax.numpy.where(adj > 0, e, zero_vec)
        attention = nn.softmax(attention)

        h_new = jax.numpy.matmul(attention, wh)

        return nn.elu(h_new)

benchmark/jax/model/gat.py
alibaba/easydist
# code modified from https://github.com/jenkspt/gpt-jax

from dataclasses import dataclass
from typing import Any, Optional

import flax.linen as nn
import jax
import jax.numpy as jnp


@dataclass(frozen=True)
class GPTConfig:
    block_size: int = 1024
    vocab_size: int = 50304
    num_layers: int = 12
    num_heads: int = 12
    num_embeds: int = 768
    dropout_rate: float = 0.0
    use_bias: bool = True
    dtype: Optional[str] = "float32"


class SelfAttention(nn.Module):

    num_heads: int
    dtype: Any = jnp.float32
    dropout_rate: float = 0.1
    deterministic: Optional[bool] = None
    use_proj_bias: bool = True

    @nn.compact
    def __call__(self, x, mask, deterministic=None):
        B, T, C = x.shape
        assert C % self.num_heads == 0
        head_dim = C // self.num_heads
        # deterministic = nn.merge_param('deterministic', self.deterministic, deterministic)

        qkv = nn.Dense(3 * C, use_bias=self.use_proj_bias, dtype=self.dtype, name='c_attn')(x)
        qkv = qkv.reshape(B, T, 3 * self.num_heads, head_dim)
        q, k, v = jnp.array_split(qkv, 3, axis=2)
        # calculate attention matrix
        scale = 1.0 / jnp.sqrt(head_dim).astype(self.dtype)
        # attn weight shape is (batch..., num_heads, q_length, kv_length)
        attn = jnp.einsum('...qhd,...khd->...hqk', q, k) * scale
        attn = jnp.where(mask, attn, jnp.finfo(self.dtype).min)
        attn = jax.nn.softmax(attn).astype(self.dtype)
        attn = nn.Dropout(self.dropout_rate, deterministic=deterministic)(attn)

        # return weighted sum over values for each query position
        x = jnp.einsum('...hqk,...khd->...qhd', attn, v).reshape(B, T, C)
        x = nn.Dense(C, use_bias=self.use_proj_bias, dtype=self.dtype, name='c_proj')(x)

        x = nn.Dropout(rate=self.dropout_rate, deterministic=deterministic)(x)
        return x


class FeedForward(nn.Module):
    config: GPTConfig

    @nn.compact
    def __call__(self, x, deterministic=None):
        B, T, C = x.shape
        x = nn.Dense(4 * C, dtype=self.config.dtype, use_bias=self.config.use_bias, name='c_fc')(x)
        x = nn.gelu(x, approximate=True)
        x = nn.Dense(C, dtype=self.config.dtype, use_bias=self.config.use_bias, name='c_proj')(x)
        x = nn.Dropout(self.config.dropout_rate, deterministic=deterministic)(x)
        return x


class GPTBlock(nn.Module):
    config: GPTConfig

    def setup(self):
        self.ln_1 = nn.LayerNorm(epsilon=1e-5,
                                 dtype=self.config.dtype,
                                 use_bias=self.config.use_bias)
        self.attn = SelfAttention(self.config.num_heads,
                                  self.config.dtype,
                                  dropout_rate=self.config.dropout_rate)
        self.ln_2 = nn.LayerNorm(epsilon=1e-5,
                                 dtype=self.config.dtype,
                                 use_bias=self.config.use_bias)
        self.mlp = FeedForward(self.config)

    def __call__(self, x, mask=None, deterministic=None):
        if mask == None:
            attn_mask = nn.make_causal_mask(jax.numpy.ones((x.shape[0], x.shape[1])), dtype=bool)
        x = x + self.attn(self.ln_1(x), attn_mask, deterministic)
        x = x + self.mlp(self.ln_2(x), deterministic)
        return x


class GPTSimple(nn.Module):
    config: GPTConfig

    @nn.compact
    def __call__(self, x, mask=None, deterministic=None):

        for i in range(self.config.num_layers):
            x = GPTBlock(self.config, name=str(i))(x, None, deterministic=deterministic)
        return x


class GPT(nn.Module):
    config: GPTConfig

    @nn.compact
    def __call__(self, idx, deterministic=None):
        B, T = idx.shape
        assert T <= self.config.block_size, f"Cannot forward sequence of length {T}, block size is only {self.block_size}"

        pos = jnp.arange(0, T)[None]
        attn_mask = nn.make_causal_mask(idx, dtype=bool)

        wte = nn.Embed(self.config.vocab_size,
                       self.config.num_embeds,
                       dtype=self.config.dtype,
                       name='wte')
        wpe = nn.Embed(self.config.block_size,
                       self.config.num_embeds,
                       dtype=self.config.dtype,
                       name='wpe')

        token_embed = wte(idx)  # [B, T, num_embeds]
        pos_embed = wpe(pos)  # [1, T, num_embeds]
        x = nn.Dropout(self.config.dropout_rate)(token_embed + pos_embed, deterministic)

        for i in range(self.config.num_layers):
            x = GPTBlock(self.config, name=str(i))(x, attn_mask, deterministic=deterministic)

        x = nn.LayerNorm(1e-5, dtype=self.config.dtype, use_bias=self.config.use_bias,
                         name='ln_f')(x)
        logits = wte.attend(x)
        return logits

    def init(self, rng):
        """
        by jitting init, traced values instead of concrete values are used
        which saves memory (since un-jitted model may not fit in memory)
        """
        tokens = jnp.zeros((2, self.config.block_size), dtype=jnp.uint16)
        params = jax.jit(super().init, static_argnums=(2, ))(rng, tokens, True)
        return params

benchmark/jax/model/gpt.py
alibaba/easydist
"""Flax implementation of ResNet V1."""

# Copyright 2023 The Flax Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

# See issue #620.
# pytype: disable=wrong-arg-count

from functools import partial
from typing import Any, Callable, Sequence, Tuple

from jax import numpy as jnp
from flax import linen as nn

ModuleDef = Any


class ResNetBlock(nn.Module):
    """ResNet block."""
    filters: int
    conv: ModuleDef
    norm: ModuleDef
    act: Callable
    strides: Tuple[int, int] = (1, 1)

    @nn.compact
    def __call__(
        self,
        x,
    ):
        residual = x
        y = self.conv(self.filters, (3, 3), self.strides)(x)
        y = self.norm()(y)
        y = self.act(y)
        y = self.conv(self.filters, (3, 3))(y)
        y = self.norm(scale_init=nn.initializers.zeros_init())(y)

        if residual.shape != y.shape:
            residual = self.conv(self.filters, (1, 1), self.strides, name='conv_proj')(residual)
            residual = self.norm(name='norm_proj')(residual)

        return self.act(residual + y)


class BottleneckResNetBlock(nn.Module):
    """Bottleneck ResNet block."""
    filters: int
    conv: ModuleDef
    norm: ModuleDef
    act: Callable
    strides: Tuple[int, int] = (1, 1)

    @nn.compact
    def __call__(self, x):
        residual = x
        y = self.conv(self.filters, (1, 1))(x)
        y = self.norm()(y)
        y = self.act(y)
        y = self.conv(self.filters, (3, 3), self.strides)(y)
        y = self.norm()(y)
        y = self.act(y)
        y = self.conv(self.filters * 4, (1, 1))(y)
        y = self.norm(scale_init=nn.initializers.zeros_init())(y)

        if residual.shape != y.shape:
            residual = self.conv(self.filters * 4, (1, 1), self.strides,
                                 name='conv_proj')(residual)
            residual = self.norm(name='norm_proj')(residual)

        return self.act(residual + y)


class ResNet(nn.Module):
    """ResNetV1."""
    stage_sizes: Sequence[int]
    block_cls: ModuleDef
    num_classes: int = 1000
    num_filters: int = 64
    dtype: Any = jnp.float32
    act: Callable = nn.relu
    conv: ModuleDef = nn.Conv

    @nn.compact
    def __call__(self, x, train: bool = True):
        conv = partial(self.conv, use_bias=False, dtype=self.dtype)
        norm = partial(nn.BatchNorm,
                       use_running_average=not train,
                       momentum=0.9,
                       epsilon=1e-5,
                       dtype=self.dtype)

        x = conv(self.num_filters, (7, 7), (2, 2), padding=[(3, 3), (3, 3)], name='conv_init')(x)
        x = norm(name='bn_init')(x)
        x = nn.relu(x)
        x = nn.max_pool(x, (3, 3), strides=(2, 2), padding='SAME')
        for i, block_size in enumerate(self.stage_sizes):
            for j in range(block_size):
                strides = (2, 2) if i > 0 and j == 0 else (1, 1)
                x = self.block_cls(self.num_filters * 2**i,
                                   strides=strides,
                                   conv=conv,
                                   norm=norm,
                                   act=self.act)(x)
        x = jnp.mean(x, axis=(1, 2))
        x = nn.Dense(self.num_classes, dtype=self.dtype)(x)
        x = jnp.asarray(x, self.dtype)
        return x


ResNet18 = partial(ResNet, stage_sizes=[2, 2, 2, 2], block_cls=ResNetBlock)

benchmark/jax/model/resnet.py
alibaba/easydist
"""The definition of wide-resnet.
Modified from https://github.com/google/flax/blob/main/examples/imagenet/models.py.
see also: https://arxiv.org/pdf/1605.07146.pdf
"""
# Copyright 2023 The Flax Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from functools import partial
from typing import Any, Callable, Sequence, Tuple

from flax import linen as nn
from flax.training import train_state, dynamic_scale as dynamic_scale_lib
import jax.numpy as jnp

ModuleDef = Any


class TrainState(train_state.TrainState):
    batch_stats: Any
    dynamic_scale: dynamic_scale_lib.DynamicScale


class ResNetBlock(nn.Module):
    """ResNet block."""
    filters: int
    conv: ModuleDef
    norm: ModuleDef
    act: Callable
    width_factor: int
    strides: Tuple[int, int] = (1, 1)

    @nn.compact
    def __call__(
        self,
        x,
    ):
        assert self.width_factor == 1

        residual = x
        y = self.conv(self.filters, (3, 3), self.strides)(x)
        y = self.norm()(y)
        y = self.act(y)
        y = self.conv(self.filters, (3, 3))(y)
        y = self.norm(scale_init=nn.initializers.zeros)(y)

        if residual.shape != y.shape:
            residual = self.conv(self.filters, (1, 1), self.strides, name='conv_proj')(residual)
            residual = self.norm(name='norm_proj')(residual)

        return self.act(residual + y)


class BottleneckResNetBlock(nn.Module):
    """Bottleneck ResNet block."""
    filters: int
    conv: ModuleDef
    norm: ModuleDef
    act: Callable
    width_factor: int
    strides: Tuple[int, int] = (1, 1)

    @nn.compact
    def __call__(self, x):
        residual = x
        y = self.conv(self.filters * self.width_factor, (1, 1))(x)
        y = self.norm()(y)
        y = self.act(y)
        y = self.conv(self.filters * self.width_factor, (3, 3), self.strides)(y)
        y = self.norm()(y)
        y = self.act(y)
        y = self.conv(self.filters * 4, (1, 1))(y)
        y = self.norm(scale_init=nn.initializers.zeros)(y)

        if residual.shape != y.shape:
            residual = self.conv(self.filters * 4, (1, 1), self.strides,
                                 name='conv_proj')(residual)
            residual = self.norm(name='norm_proj')(residual)

        return self.act(residual + y)


class ResNet(nn.Module):
    """ResNetV1."""
    stage_sizes: Sequence[int]
    block_cls: ModuleDef
    num_filters: int = 64
    width_factor: int = 1
    num_classes: int = 1000
    dtype: Any = jnp.float32
    act: Callable = nn.relu

    @nn.compact
    def __call__(self, x, train: bool = True):
        conv = partial(nn.Conv, use_bias=False, dtype=self.dtype)
        norm = partial(nn.BatchNorm,
                       use_running_average=not train,
                       momentum=0.9,
                       epsilon=1e-5,
                       dtype=self.dtype)

        x = conv(self.num_filters, (7, 7), (2, 2), padding=[(3, 3), (3, 3)], name='conv_init')(x)
        x = norm(name='bn_init')(x)
        x = nn.relu(x)
        x = nn.max_pool(x, (3, 3), strides=(2, 2), padding='SAME')
        for i, block_size in enumerate(self.stage_sizes):
            for j in range(block_size):
                strides = (2, 2) if i > 0 and j == 0 else (1, 1)
                x = self.block_cls(self.num_filters * 2**i,
                                   strides=strides,
                                   conv=conv,
                                   norm=norm,
                                   width_factor=self.width_factor,
                                   act=self.act)(x)
        x = jnp.mean(x, axis=(1, 2))
        x = nn.Dense(self.num_classes, dtype=self.dtype)(x)
        x = jnp.asarray(x, self.dtype)
        return x


model_configs = {
    0: {
        "stage_sizes": [],
        "block_cls": ResNetBlock
    },
    18: {
        "stage_sizes": [2, 2, 2, 2],
        "block_cls": ResNetBlock
    },
    34: {
        "stage_sizes": [3, 4, 6, 3],
        "block_cls": ResNetBlock
    },
    50: {
        "stage_sizes": [3, 4, 6, 3],
        "block_cls": BottleneckResNetBlock
    },
    101: {
        "stage_sizes": [3, 4, 23, 3],
        "block_cls": BottleneckResNetBlock
    },
    152: {
        "stage_sizes": [3, 8, 36, 3],
        "block_cls": BottleneckResNetBlock
    },
    200: {
        "stage_sizes": [3, 24, 36, 3],
        "block_cls": BottleneckResNetBlock
    }
}


def get_wide_resnet(num_layers, width_factor, num_filters, num_classes, dtype):
    model_config = model_configs[num_layers]
    model_config["width_factor"] = width_factor
    model_config["num_filters"] = num_filters
    model_config["num_classes"] = num_classes
    model_config["dtype"] = dtype

    return ResNet(**model_config)


def resnet18():
    return ResNet(stage_sizes=[2, 2, 2, 2], block_cls=ResNetBlock)


def resnet34():
    return ResNet(stage_sizes=[3, 4, 6, 3], block_cls=ResNetBlock)


def resnet50():
    return ResNet(stage_sizes=[3, 4, 6, 3], block_cls=BottleneckResNetBlock)


def wresnet50():
    return ResNet(block_cls=BottleneckResNetBlock, stage_sizes=[3, 4, 6, 3], width_factor=2)


def wresnet101():
    return ResNet(block_cls=BottleneckResNetBlock, stage_sizes=[3, 4, 23, 3], width_factor=2)

benchmark/jax/model/wresnet.py
alibaba/easydist
# torchrun --nproc_per_node 2 --master_port 26543 ./benchmark/bench_torch.py

import logging
import os
import sys
from functools import partial

import numpy
import torch
import torch.optim as optim
import torch.utils._pytree as pytree
from torch.distributed._tensor import DeviceMesh
from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
from torch.nn.parallel import DistributedDataParallel as DDP

from easydist.torch.experimental.api import easydist_compile
from easydist.torch import (enable_transform, get_input_strategy, compile, set_device_mesh,
                            shard_module)
from easydist.utils.testing import TorchMockDeviceMesh
from easydist import easydist_setup
from easydist.utils.timer import EDTimer

sys.path.append(os.path.dirname(os.path.abspath(__file__)))
from benchmark.torch.model import GPT, GATLayer, wresnet50
from benchmark.bench_case import GPTCase, ResNetCase, GATCase

logging.basicConfig(format='%(asctime)s - %(levelname)s - %(message)s',
                    datefmt='%m/%d %H:%M:%S',
                    level=logging.INFO)


def get_gpt_case(device="cuda"):

    case = GPTCase()
    model = GPT(depth=case.num_layers, dim=case.hidden_dim, num_heads=case.num_heads)
    data_in = torch.ones(case.batch_size, case.seq_size, case.hidden_dim)

    model.device = torch.device(device)
    return model.to(device=device), data_in.to(device=device)


def get_resnet_case(device="cuda"):
    case = ResNetCase()
    model = wresnet50()
    data_in = torch.ones(case.batch_size, 3, 224, 224)

    model.device = torch.device(device)
    return model.to(device=device), data_in.to(device=device)


def get_gat_case(device="cuda"):
    case = GATCase()
    model = GATLayer(case.in_feature, case.out_feature)
    data_in = torch.ones(case.num_node, case.in_feature)
    adj = torch.ones(case.num_node, case.num_node)

    model.device = torch.device(device)
    return model.to(device=device), [data_in.to(device=device), adj.to(device=device)]


def bench_ddp(model, data_in):

    if not isinstance(data_in, list):
        data_in = [data_in]

    world_size = torch.distributed.get_world_size()
    for i in range(len(data_in)):
        data_in[i] = torch.chunk(data_in[i], world_size)[0]

    ddp_model = DDP(model)
    optimizer = optim.SGD(ddp_model.parameters(), lr=0.001)

    def train_step():
        optimizer.zero_grad()
        out = ddp_model(*data_in)
        out_grad = torch.ones_like(out)
        out.backward(out_grad)
        optimizer.step()

    torch.cuda.reset_peak_memory_stats()

    timer = EDTimer(train_step, in_ms=False)

    elaps_time = timer.time()
    peak_memory = torch.cuda.max_memory_allocated()

    print(f"Memory: {peak_memory / 1024 / 1024 / 1024} GB")
    print(f"Time: {elaps_time}")


def bench_fsdp(model, data_in):

    if not isinstance(data_in, list):
        data_in = [data_in]

    if not isinstance(model, GATLayer):
        world_size = torch.distributed.get_world_size()
        for i in range(len(data_in)):
            data_in[i] = torch.chunk(data_in[i], world_size)[0]

    ddp_model = FSDP(model)
    optimizer = optim.SGD(ddp_model.parameters(), lr=0.001)

    def train_step():
        optimizer.zero_grad()
        out = ddp_model(*data_in)
        out_grad = torch.ones_like(out)
        out.backward(out_grad)
        optimizer.step()

    torch.cuda.reset_peak_memory_stats()

    timer = EDTimer(train_step, in_ms=False)

    elaps_time = timer.time()
    peak_memory = torch.cuda.max_memory_allocated()

    print(f"Memory: {peak_memory / 1024 / 1024 / 1024} GB")
    print(f"Time: {elaps_time}")


def to_meta(node_output):
    if type(node_output) is torch.Tensor:
        return node_output.to(device="meta")
    elif type(node_output) is torch.nn.parameter.Parameter:
        return node_output.to(device="meta")
    else:
        return node_output


def bench_easydist(model, data_in):
    mock_mesh = TorchMockDeviceMesh(1, 2, debug_only=True)
    set_device_mesh(mock_mesh)

    if not isinstance(data_in, list):
        data_in = [data_in]

    if model.device == torch.device("meta"):
        data_in = pytree.tree_map(to_meta, data_in)

    optimizer = optim.SGD(model.parameters(), lr=0.001)

    @easydist_compile()
    def train_step(model, optimizer, data_in):
        output_ = model(*data_in)
        output_grad = torch.ones_like(output_)
        output_.backward(output_grad)
        optimizer.step()
        optimizer.zero_grad()
        return output_

    torch.cuda.empty_cache()
    torch.cuda.reset_peak_memory_stats()

    train_step_partial = partial(train_step, model, optimizer, data_in)

    timer = EDTimer(train_step_partial, in_ms=False)

    elaps_time = timer.time()
    peak_memory = torch.cuda.max_memory_allocated()

    print(f"Memory: {peak_memory / 1024 / 1024 / 1024} GB")
    print(f"Time: {elaps_time}")


def bench_easydist_old(model, data_in):
    world_size = torch.distributed.get_world_size()

    mesh_shape = numpy.array(range(world_size)).reshape(1, -1)
    mesh = DeviceMesh("cuda", mesh_shape.tolist())

    mock_mesh = TorchMockDeviceMesh(*(mesh_shape.shape))
    set_device_mesh(mock_mesh)

    enable_transform()

    compiled_module = compile(model)

    if not isinstance(data_in, list):
        data_in = [data_in]

    if model.device == torch.device("meta"):
        data_in = pytree.tree_map(to_meta, data_in)

    out = compiled_module(*data_in)
    try:
        out_grad = torch.ones_like(out)
        out.backward(out_grad)
    except:
        pass

    set_device_mesh(mesh)
    data_in = shard_module(compiled_module.orig_module, data_in, get_input_strategy())

    # handle causal_mask for GPT
    if isinstance(compiled_module.orig_module, GPT):
        for idx in range(len(compiled_module.orig_module.blocks)):

            tensor_data = compiled_module.orig_module.blocks[idx].attn.core_attention.causal_mask

            if tensor_data.device == torch.device("meta"):
                seq_len = tensor_data.shape[-1]
                tensor_data = torch.tril(
                    torch.ones((seq_len, seq_len), dtype=torch.uint8,
                               device="cuda")).view(1, 1, seq_len, seq_len).bool()

            compiled_module.orig_module.blocks[idx].attn.core_attention.causal_mask = tensor_data

    optimizer = optim.SGD(compiled_module.orig_module.parameters(), lr=0.001)

    def train_step():
        optimizer.zero_grad()
        out = compiled_module(*data_in)
        out_grad = torch.ones_like(out)
        out.backward(out_grad)
        optimizer.step()

    torch.cuda.empty_cache()
    torch.cuda.reset_peak_memory_stats()

    timer = EDTimer(train_step, in_ms=False)

    elaps_time = timer.time()
    peak_memory = torch.cuda.max_memory_allocated()

    print(f"Memory: {peak_memory / 1024 / 1024 / 1024} GB")
    print(f"Time: {elaps_time}")


def main():
    # setup easydist
    easydist_setup(backend="torch", device="cuda")

    # setup distributed
    torch.distributed.init_process_group(backend="nccl")
    local_rank = int(os.environ["LOCAL_RANK"])
    torch.cuda.set_device(local_rank)

    model, data_in = get_gpt_case(device="meta")

    bench_easydist(model, data_in)


if __name__ == '__main__':
    main()

benchmark/torch/bench_torch.py
alibaba/easydist
import os
import sys

import torch
import torch.optim as optim
from torch.nn.parallel import DistributedDataParallel as DDP
from fairscale.nn.model_parallel.initialize import initialize_model_parallel
from fairscale.nn.model_parallel import get_data_parallel_group

from easydist.utils.timer import EDTimer
from easydist import easydist_setup

sys.path.append(os.path.dirname(os.path.abspath(__file__)))
from benchmark.bench_case import GPTCase
from benchmark.torch.model.gpt_tp import GPT


def get_gpt_case(cuda=True):

    case = GPTCase()
    model = GPT(depth=case.num_layers, dim=case.hidden_dim, num_heads=case.num_heads)
    data_in = torch.ones(case.batch_size, case.seq_size, case.hidden_dim)

    if cuda:
        return model.cuda(), data_in.cuda()

    return model, data_in


def bench_tp(model, data_in):

    ddp_model = DDP(model, process_group=get_data_parallel_group())
    optimizer = optim.SGD(ddp_model.parameters(), lr=0.001)

    def train_step():
        optimizer.zero_grad()
        out = ddp_model(data_in)
        out_grad = torch.ones_like(out)
        out.backward(out_grad)
        optimizer.step()

    torch.cuda.reset_peak_memory_stats()

    timer = EDTimer(train_step, in_ms=False)

    elaps_time = timer.time()
    peak_memory = torch.cuda.max_memory_allocated()

    print(f"Memory: {peak_memory / 1024 / 1024 / 1024} GB")
    print(f"Time: {elaps_time}")


def main():
    easydist_setup(backend="torch", device="cuda")
    # setup distributed
    torch.distributed.init_process_group(backend="nccl")
    local_rank = int(os.environ["LOCAL_RANK"])
    world_size = torch.distributed.get_world_size()
    initialize_model_parallel(world_size)
    torch.cuda.set_device(local_rank)

    model, data_in = get_gpt_case(cuda=True)

    bench_tp(model, data_in)


if __name__ == '__main__':
    main()

benchmark/torch/bench_torch_tp.py
alibaba/easydist
from .gat import GATLayer
from .gpt import GPT, GPTLayer, FeedForward, SelfAttention
from .wresnet import resnet18, resnet34, resnet50, wresnet50, wresnet101
from .llama import LLAMA, LLAMAConfig

__all__ = [
    'GATLayer', 'GPT', 'GPTLayer', 'FeedForward', 'SelfAttention', 'wresnet50', 'wresnet101',
    'resnet18', 'resnet34', 'resnet50', 'LLAMA', 'LLAMAConfig'
]

benchmark/torch/model/__init__.py
alibaba/easydist
# code modified from https://github.com/Diego999/pyGAT

import torch
import torch.nn as nn
import torch.nn.functional as F


class GATLayer(nn.Module):

    def __init__(self, in_features, out_features):
        super(GATLayer, self).__init__()
        self.in_features = in_features
        self.out_features = out_features

        self.linear_w = nn.Linear(in_features=in_features, out_features=out_features, bias=None)

        self.linear_a_1 = nn.Linear(in_features=out_features, out_features=1, bias=None)
        self.linear_a_2 = nn.Linear(in_features=out_features, out_features=1, bias=None)

        self.leakyrelu = nn.LeakyReLU()

        self.softmax = nn.Softmax(dim=-1)

    def forward(self, h, adj):
        wh = self.linear_w(h)
        wh1 = self.linear_a_1(wh)
        wh2 = self.linear_a_2(wh)

        e = self.leakyrelu(wh1 + wh2.T)
        zero_vec = -10e10 * torch.ones(*e.shape, device=e.device, dtype=e.dtype)
        attention = torch.where(adj > 0, e, zero_vec)
        attention = self.softmax(attention)

        h_new = torch.matmul(attention, wh)

        return F.elu(h_new)

benchmark/torch/model/gat.py
alibaba/easydist
import math

import torch
import torch.nn as nn
import torch.nn.functional as F


class AttentionCore(nn.Module):

    def __init__(self, attention_head_size: int, attention_dropout: float) -> None:
        super().__init__()
        self.attention_head_size = attention_head_size
        self.attention_dropout = nn.Dropout(attention_dropout)

        self.softmax = nn.Softmax(dim=-1)

        # (avoid the const tensor init when forward)
        # self.register_buffer('causal_mask', None)
        self.causal_mask = None
        self.where_const = -1e4

    def forward(self, q, k, v):
        x = torch.matmul(q, k.transpose(-1, -2))
        x = x / math.sqrt(self.attention_head_size)

        # (avoid the const tensor init when forward)
        if self.causal_mask is None:
            q_len, k_len = q.size(-2), k.size(-2)
            self.causal_mask = torch.tril(
                torch.ones((q_len, k_len), dtype=torch.uint8,
                           device=q.device)).view(1, 1, q_len, k_len).bool()
        x = torch.where(self.causal_mask, x, self.where_const)
        x = self.softmax(x)
        x = self.attention_dropout(x)

        x = torch.matmul(x, v)
        x = x.transpose(1, 2)
        new_context_layer_shape = x.size()[:-2] + (-1, )
        x = x.reshape(new_context_layer_shape)

        return x


class FeedForward(nn.Module):

    def __init__(self, hidden_size, ratio=4) -> None:
        super().__init__()
        self.hidden_size = hidden_size
        self.ffn_hidden_size = hidden_size * ratio

        self.dense_h_to_4h = nn.Linear(self.hidden_size, self.ffn_hidden_size)
        self.dense_4h_to_h = nn.Linear(self.ffn_hidden_size, self.hidden_size)

        self.activation_func = F.gelu

    def forward(self, hidden_states):
        intermediate_parallel = self.dense_h_to_4h(hidden_states)
        intermediate_parallel = self.activation_func(intermediate_parallel)
        output = self.dense_4h_to_h(intermediate_parallel)
        return output


class SelfAttention(nn.Module):

    def __init__(self,
                 dim: int,
                 num_heads: int,
                 attention_dropout: float = 0.,
                 dropout: float = 0.) -> None:
        super().__init__()

        self.attention_head_size = dim // num_heads
        self.query = nn.Linear(dim, dim)
        self.key = nn.Linear(dim, dim)
        self.value = nn.Linear(dim, dim)

        self.dense = nn.Linear(dim, dim)

        self.dropout = nn.Dropout(dropout)

        self.core_attention = AttentionCore(self.attention_head_size, attention_dropout)

    def forward(self, x):

        q = self.query(x)
        k = self.key(x)
        v = self.value(x)

        all_head_size = q.shape[-1]
        num_attention_heads = all_head_size // self.attention_head_size

        local_attention_heads = num_attention_heads
        new_qkv_shape = q.shape[:-1] + \
            (local_attention_heads, self.attention_head_size)
        q = q.view(new_qkv_shape)
        k = k.view(new_qkv_shape)
        v = v.view(new_qkv_shape)

        q = q.permute((0, 2, 1, 3))
        k = k.permute((0, 2, 1, 3))
        v = v.permute((0, 2, 1, 3))

        x = self.core_attention(q, k, v)

        x = self.dense(x)
        x = self.dropout(x)

        return x


class GPTLayer(nn.Module):

    def __init__(self,
                 dim: int,
                 num_heads: int,
                 mlp_ratio: int = 4,
                 attention_dropout: float = 0.,
                 dropout: float = 0.,
                 dtype: torch.dtype = None):
        super().__init__()
        self.norm1 = nn.LayerNorm(normalized_shape=dim, eps=1e-6, dtype=dtype)
        self.attn = SelfAttention(dim=dim,
                                  num_heads=num_heads,
                                  attention_dropout=attention_dropout,
                                  dropout=dropout)
        self.norm2 = nn.LayerNorm(normalized_shape=dim, eps=1e-6, dtype=dtype)
        self.mlp = FeedForward(hidden_size=dim, ratio=mlp_ratio)

    def forward(self, x):
        x = x + self.attn(self.norm1(x))
        x = x + self.mlp(self.norm2(x))
        return x


class GPT(nn.Module):

    def __init__(self,
                 depth: int,
                 dim: int,
                 num_heads: int,
                 mlp_ratio: int = 4,
                 attention_dropout: float = 0.,
                 dropout: float = 0.,
                 dtype: torch.dtype = None):
        super().__init__()
        self.blocks = nn.ModuleList([
            GPTLayer(
                dim=dim,
                num_heads=num_heads,
                mlp_ratio=mlp_ratio,
                attention_dropout=attention_dropout,
                dropout=dropout,
                dtype=dtype,
            ) for _ in range(depth)
        ])

    def forward(self, x):
        for block in self.blocks:
            x = block(x)
        return x

benchmark/torch/model/gpt.py
alibaba/easydist
import math
import os

import torch
import torch.nn as nn
import torch.nn.functional as F

import fairscale.nn.model_parallel.initialize as fs_init
from fairscale.nn.model_parallel.layers import ColumnParallelLinear, RowParallelLinear


class AttentionCore(nn.Module):

    def __init__(self, attention_head_size: int, attention_dropout: float) -> None:
        super().__init__()
        self.attention_head_size = attention_head_size
        self.attention_dropout = nn.Dropout(attention_dropout)

        self.softmax = nn.Softmax(dim=-1)

        # (avoid the const tensor init when forward)
        self.causal_mask = None
        self.where_const = -1e4

    def forward(self, q, k, v):
        x = torch.matmul(q, k.transpose(-1, -2))
        x = x / math.sqrt(self.attention_head_size)

        # (avoid the const tensor init when forward)
        if self.causal_mask is None:
            q_len, k_len = q.size(-2), k.size(-2)
            meta_dev = os.environ.get("EASYDIST_DEVICE", "cuda")
            self.causal_mask = torch.tril(
                torch.ones((q_len, k_len), dtype=torch.uint8,
                           device=meta_dev)).view(1, 1, q_len, k_len).bool()
        x = torch.where(self.causal_mask, x, self.where_const)
        x = self.softmax(x)
        x = self.attention_dropout(x)

        x = torch.matmul(x, v)
        x = x.transpose(1, 2)
        new_context_layer_shape = x.size()[:-2] + (-1, )
        x = x.reshape(new_context_layer_shape)

        return x


class FeedForward(nn.Module):

    def __init__(self, hidden_size, ratio=4) -> None:
        super().__init__()
        self.hidden_size = hidden_size
        self.ffn_hidden_size = hidden_size * ratio

        self.dense_h_to_4h = ColumnParallelLinear(self.hidden_size,
                                                  self.ffn_hidden_size,
                                                  gather_output=False)
        self.dense_4h_to_h = RowParallelLinear(self.ffn_hidden_size,
                                               self.hidden_size,
                                               input_is_parallel=True)

        self.activation_func = F.gelu

    def forward(self, hidden_states):
        intermediate_parallel = self.dense_h_to_4h(hidden_states)
        intermediate_parallel = self.activation_func(intermediate_parallel)
        output = self.dense_4h_to_h(intermediate_parallel)
        return output


class SelfAttention(nn.Module):

    def __init__(self,
                 dim: int,
                 num_heads: int,
                 attention_dropout: float = 0.,
                 dropout: float = 0.) -> None:
        super().__init__()

        self.attention_head_size = dim // num_heads
        self.query = ColumnParallelLinear(dim, dim, gather_output=False)
        self.key = ColumnParallelLinear(dim, dim, gather_output=False)
        self.value = ColumnParallelLinear(dim, dim, gather_output=False)

        self.dense = RowParallelLinear(dim, dim, input_is_parallel=True)

        self.dropout = nn.Dropout(dropout)

        self.core_attention = AttentionCore(self.attention_head_size, attention_dropout)

    def forward(self, x):

        q = self.query(x)
        k = self.key(x)
        v = self.value(x)

        all_head_size = q.shape[-1]
        num_attention_heads = all_head_size // self.attention_head_size

        new_qkv_shape = q.shape[:-1] + \
            (num_attention_heads, self.attention_head_size)
        q = q.view(new_qkv_shape)
        k = k.view(new_qkv_shape)
        v = v.view(new_qkv_shape)

        q = q.permute((0, 2, 1, 3))
        k = k.permute((0, 2, 1, 3))
        v = v.permute((0, 2, 1, 3))

        x = self.core_attention(q, k, v)

        x = self.dense(x)
        x = self.dropout(x)

        return x


class GPTLayer(nn.Module):

    def __init__(self,
                 dim: int,
                 num_heads: int,
                 mlp_ratio: int = 4,
                 attention_dropout: float = 0.,
                 dropout: float = 0.,
                 dtype: torch.dtype = None):
        super().__init__()
        self.norm1 = nn.LayerNorm(normalized_shape=dim, eps=1e-6, dtype=dtype)
        self.attn = SelfAttention(dim=dim,
                                  num_heads=num_heads,
                                  attention_dropout=attention_dropout,
                                  dropout=dropout)
        self.norm2 = nn.LayerNorm(normalized_shape=dim, eps=1e-6, dtype=dtype)
        self.mlp = FeedForward(hidden_size=dim, ratio=mlp_ratio)

    def forward(self, x):
        x = x + self.attn(self.norm1(x))
        x = x + self.mlp(self.norm2(x))
        return x


class GPT(nn.Module):

    def __init__(self,
                 depth: int,
                 dim: int,
                 num_heads: int,
                 mlp_ratio: int = 4,
                 attention_dropout: float = 0.,
                 dropout: float = 0.,
                 dtype: torch.dtype = None):
        super().__init__()
        self.blocks = nn.ModuleList([
            GPTLayer(
                dim=dim,
                num_heads=num_heads,
                mlp_ratio=mlp_ratio,
                attention_dropout=attention_dropout,
                dropout=dropout,
                dtype=dtype,
            ) for _ in range(depth)
        ])

    def forward(self, x):
        for block in self.blocks:
            x = block(x)
        return x

benchmark/torch/model/gpt_tp.py
alibaba/easydist
# code modified from https://github.com/pytorch/vision/blob/main/torchvision/models/resnet.py

import torch
from torch import Tensor
import torch.nn as nn
from typing import Type, Any, Callable, Union, List, Optional


def conv3x3(in_planes: int,
            out_planes: int,
            stride: int = 1,
            groups: int = 1,
            dilation: int = 1) -> nn.Conv2d:
    """3x3 convolution with padding"""
    return nn.Conv2d(in_planes,
                     out_planes,
                     kernel_size=3,
                     stride=stride,
                     padding=dilation,
                     groups=groups,
                     bias=False,
                     dilation=dilation)


def conv1x1(in_planes: int, out_planes: int, stride: int = 1) -> nn.Conv2d:
    """1x1 convolution"""
    return nn.Conv2d(in_planes, out_planes, kernel_size=1, stride=stride, bias=False)


class BasicBlock(nn.Module):
    expansion: int = 1

    def __init__(self,
                 inplanes: int,
                 planes: int,
                 stride: int = 1,
                 downsample: Optional[nn.Module] = None,
                 groups: int = 1,
                 base_width: int = 64,
                 dilation: int = 1,
                 norm_layer: Optional[Callable[..., nn.Module]] = None) -> None:
        super(BasicBlock, self).__init__()
        if norm_layer is None:
            norm_layer = nn.BatchNorm2d
        if groups != 1 or base_width != 64:
            raise ValueError('BasicBlock only supports groups=1 and base_width=64')
        if dilation > 1:
            raise NotImplementedError("Dilation > 1 not supported in BasicBlock")
        # Both self.conv1 and self.downsample layers downsample the input when stride != 1
        self.conv1 = conv3x3(inplanes, planes, stride)
        self.bn1 = norm_layer(planes)
        self.relu = nn.ReLU(inplace=True)
        self.conv2 = conv3x3(planes, planes)
        self.bn2 = norm_layer(planes)
        self.downsample = downsample
        self.stride = stride

    def forward(self, x: Tensor) -> Tensor:
        identity = x

        out = self.conv1(x)
        out = self.bn1(out)
        out = self.relu(out)

        out = self.conv2(out)
        out = self.bn2(out)

        if self.downsample is not None:
            identity = self.downsample(x)

        out += identity
        out = self.relu(out)

        return out


class Bottleneck(nn.Module):
    # Bottleneck in torchvision places the stride for downsampling at 3x3 convolution(self.conv2)
    # while original implementation places the stride at the first 1x1 convolution(self.conv1)
    # according to "Deep residual learning for image recognition"https://arxiv.org/abs/1512.03385.
    # This variant is also known as ResNet V1.5 and improves accuracy according to
    # https://ngc.nvidia.com/catalog/model-scripts/nvidia:resnet_50_v1_5_for_pytorch.

    expansion: int = 4

    def __init__(self,
                 inplanes: int,
                 planes: int,
                 stride: int = 1,
                 downsample: Optional[nn.Module] = None,
                 groups: int = 1,
                 base_width: int = 64,
                 dilation: int = 1,
                 norm_layer: Optional[Callable[..., nn.Module]] = None) -> None:
        super(Bottleneck, self).__init__()
        if norm_layer is None:
            norm_layer = nn.BatchNorm2d
        width = int(planes * (base_width / 64.)) * groups
        # Both self.conv2 and self.downsample layers downsample the input when stride != 1
        self.conv1 = conv1x1(inplanes, width)
        self.bn1 = norm_layer(width)
        self.conv2 = conv3x3(width, width, stride, groups, dilation)
        self.bn2 = norm_layer(width)
        self.conv3 = conv1x1(width, planes * self.expansion)
        self.bn3 = norm_layer(planes * self.expansion)
        self.relu = nn.ReLU(inplace=True)
        self.downsample = downsample
        self.stride = stride

    def forward(self, x: Tensor) -> Tensor:
        identity = x

        out = self.conv1(x)
        out = self.bn1(out)
        out = self.relu(out)

        out = self.conv2(out)
        out = self.bn2(out)
        out = self.relu(out)

        out = self.conv3(out)
        out = self.bn3(out)

        if self.downsample is not None:
            identity = self.downsample(x)

        out += identity
        out = self.relu(out)

        return out


class ResNet(nn.Module):

    def __init__(self,
                 block: Type[Union[BasicBlock, Bottleneck]],
                 layers: List[int],
                 num_classes: int = 1000,
                 num_filters: int = 64,
                 zero_init_residual: bool = False,
                 groups: int = 1,
                 width_per_group: int = 64,
                 replace_stride_with_dilation: Optional[List[bool]] = None,
                 norm_layer: Optional[Callable[..., nn.Module]] = None) -> None:
        super(ResNet, self).__init__()
        if norm_layer is None:
            norm_layer = nn.BatchNorm2d
        self._norm_layer = norm_layer

        self.inplanes = num_filters
        self.dilation = 1
        if replace_stride_with_dilation is None:
            # each element in the tuple indicates if we should replace
            # the 2x2 stride with a dilated convolution instead
            replace_stride_with_dilation = [False, False, False]
        if len(replace_stride_with_dilation) != 3:
            raise ValueError("replace_stride_with_dilation should be None "
                             "or a 3-element tuple, got {}".format(replace_stride_with_dilation))
        self.groups = groups
        self.base_width = width_per_group
        self.conv1 = nn.Conv2d(3, self.inplanes, kernel_size=7, stride=2, padding=3, bias=False)
        self.bn1 = norm_layer(self.inplanes)
        self.relu = nn.ReLU(inplace=True)
        self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
        self.layer1 = self._make_layer(block, 64, layers[0])
        self.layer2 = self._make_layer(block,
                                       128,
                                       layers[1],
                                       stride=2,
                                       dilate=replace_stride_with_dilation[0])
        self.layer3 = self._make_layer(block,
                                       256,
                                       layers[2],
                                       stride=2,
                                       dilate=replace_stride_with_dilation[1])
        self.layer4 = self._make_layer(block,
                                       512,
                                       layers[3],
                                       stride=2,
                                       dilate=replace_stride_with_dilation[2])
        self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
        self.fc = nn.Linear(512 * block.expansion, num_classes)

        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
            elif isinstance(m, (nn.BatchNorm2d, nn.GroupNorm)):
                nn.init.constant_(m.weight, 1)
                nn.init.constant_(m.bias, 0)

        # Zero-initialize the last BN in each residual branch,
        # so that the residual branch starts with zeros, and each residual block behaves like an identity.
        # This improves the model by 0.2~0.3% according to https://arxiv.org/abs/1706.02677
        if zero_init_residual:
            for m in self.modules():
                if isinstance(m, Bottleneck):
                    nn.init.constant_(m.bn3.weight, 0)  # type: ignore[arg-type]
                elif isinstance(m, BasicBlock):
                    nn.init.constant_(m.bn2.weight, 0)  # type: ignore[arg-type]

    def _make_layer(self,
                    block: Type[Union[BasicBlock, Bottleneck]],
                    planes: int,
                    blocks: int,
                    stride: int = 1,
                    dilate: bool = False) -> nn.Sequential:
        norm_layer = self._norm_layer
        downsample = None
        previous_dilation = self.dilation
        if dilate:
            self.dilation *= stride
            stride = 1
        if stride != 1 or self.inplanes != planes * block.expansion:
            downsample = nn.Sequential(
                conv1x1(self.inplanes, planes * block.expansion, stride),
                norm_layer(planes * block.expansion),
            )

        layers = []
        layers.append(
            block(self.inplanes, planes, stride, downsample, self.groups, self.base_width,
                  previous_dilation, norm_layer))
        self.inplanes = planes * block.expansion
        for _ in range(1, blocks):
            layers.append(
                block(self.inplanes,
                      planes,
                      groups=self.groups,
                      base_width=self.base_width,
                      dilation=self.dilation,
                      norm_layer=norm_layer))

        return nn.Sequential(*layers)

    def _forward_impl(self, x: Tensor) -> Tensor:
        # See note [TorchScript super()]
        x = self.conv1(x)
        x = self.bn1(x)
        x = self.relu(x)
        x = self.maxpool(x)

        x = self.layer1(x)
        x = self.layer2(x)
        x = self.layer3(x)
        x = self.layer4(x)

        x = self.avgpool(x)
        x = torch.flatten(x, 1)
        x = self.fc(x)

        return x

    def forward(self, x: Tensor) -> Tensor:
        return self._forward_impl(x)


def resnet18():
    return ResNet(BasicBlock, [2, 2, 2, 2])


def resnet34():
    return ResNet(BasicBlock, [3, 4, 6, 3])


def resnet50():
    return ResNet(Bottleneck, [3, 4, 6, 3])


def wresnet50():
    return ResNet(Bottleneck, [3, 4, 6, 3], width_per_group=64 * 2)


def wresnet101():
    return ResNet(Bottleneck, [3, 4, 23, 3], width_per_group=64 * 2)

benchmark/torch/model/wresnet.py
alibaba/easydist
# Copyright (c) 2023, Alibaba Group;
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================

import logging

from . import platform
import easydist.config as mdconfig


def easydist_setup(backend, device="cpu", allow_tf32=True):
    mdconfig.easydist_device = device

    logging.basicConfig(format='%(asctime)s - %(levelname)s - %(name)s - %(message)s',
                        datefmt='%m/%d %H:%M:%S',
                        level=mdconfig.log_level)

    if backend == "jax":
        from .jax import easydist_setup_jax
        easydist_setup_jax(device, allow_tf32)

        logging.getLogger("jax._src").setLevel(logging.INFO)
    elif backend == "torch":
        from .torch import easydist_setup_torch
        easydist_setup_torch(device, allow_tf32)

        logging.getLogger("torch._subclasses.fake_tensor").setLevel(logging.INFO)

    platform.init_backend(backend)

easydist/__init__.py
alibaba/easydist
# Copyright (c) 2023, Alibaba Group;
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================

from .solver import AutoFlowSolver

__all__ = ["AutoFlowSolver"]

easydist/autoflow/__init__.py
alibaba/easydist
# Copyright (c) 2023, Alibaba Group;
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================

import logging
import random
from pprint import pprint
from typing import List

import easydist.config as mdconfig
import mip
from easydist.metashard.metair import (SPMD, MetaGraph, MetaNode, MetaNodeCluster, MetaVar,
                                       NodeSPMDStrategy, VarSPMDStrategy, VarSPMDStrategyGroup,
                                       ClusterStrategyPool)

logger = logging.getLogger(__name__)


def shuffle_list(*ls):
    l = list(zip(*ls))

    random.shuffle(l)
    return zip(*l)


def get_idx_in_var_list(var: MetaVar, var_list: List[MetaVar]):
    var_list = [v.name if v else None for v in var_list]
    if var.name in var_list:
        return var_list.index(var.name)
    return None


def calculate_resharding_cost(var: MetaVar, strategy_in: VarSPMDStrategy,
                              strategy_out: VarSPMDStrategy, device_mesh):
    var_size = var.get_var_size()

    all_gather = lambda x, ndevices: x * (ndevices - 1) / ndevices
    all_reduce = lambda x, ndevices: 2 * x * (ndevices - 1) / ndevices

    def all_to_all(x, ndevices):
        factor = 1.
        if ndevices > 2:
            factor = mdconfig.all_to_all_punish_factor
        return factor * x * (ndevices - 1) / ndevices / ndevices

    resharding_cost = 0

    s1_in, s2_in = strategy_in[0], strategy_in[1]
    s1_out, s2_out = strategy_out[0], strategy_out[1]

    if device_mesh[0] > 1:
        message_size = var_size
        if s2_in.is_shard():
            message_size /= device_mesh[1]
        if s1_in.is_shard():
            if s1_out.is_shard():
                if s1_in.args != s1_out.args:
                    resharding_cost += all_to_all(message_size, device_mesh[0])
            else:
                resharding_cost += all_gather(message_size, device_mesh[0])
        elif s1_in.is_partial():
            resharding_cost += all_reduce(message_size, device_mesh[0])

    if device_mesh[1] > 1:
        message_size = var_size
        if s1_in.is_shard():
            message_size /= device_mesh[0]
        if s2_in.is_shard():
            if s2_out.is_shard():
                if s2_in.args != s2_out.args:
                    resharding_cost += all_to_all(message_size, device_mesh[1])
            else:
                resharding_cost += all_gather(message_size, device_mesh[1])
        elif s2_in.is_partial():
            resharding_cost += all_reduce(message_size, device_mesh[1])

    return resharding_cost


def calculate_memory_cost(var: MetaVar, strategy_in: VarSPMDStrategy,
                          strategy_out: VarSPMDStrategy, device_mesh):
    var_size = var.get_var_size()

    memory_cost = 0

    # FIXME: only use out_strategy here for shard_size is ok?
    for strategy in [strategy_in, strategy_out]:
        s1, s2 = strategy[0], strategy[1]
        shard_size = 1
        if s1.is_shard():
            shard_size *= device_mesh[0]
        if s2.is_shard():
            shard_size *= device_mesh[1]

        memory_cost += var_size // shard_size

    return memory_cost


def gen_comm_cost_matrix(
        var: MetaVar,
        up_strategy,  # list of VarSPMDStrategy
        down_strategy,  # list of VarSPMDStrategy
        device_mesh):
    comm_matrix = [[0 for _ in range(len(down_strategy))] for _ in range(len(up_strategy))]

    for i in range(len(up_strategy)):
        for j in range(len(down_strategy)):
            var_up_strategy = up_strategy[i]
            var_down_strategy = down_strategy[j]
            comm_matrix[i][j] = calculate_resharding_cost(var, var_up_strategy, var_down_strategy,
                                                          device_mesh)

    return comm_matrix


# Lansong(TODO) generate_comm_matrix is deprecated, use gen_comm_cost_matrix instead
def generate_comm_matrix(
        var: MetaVar,
        up_strategy,  # list of NodeSPMDStrategy
        down_strategy,  # list of NodeSPMDStrategy
        idx_for_up,
        idx_for_down,
        device_mesh):
    comm_matrix = [[0 for _ in range(len(down_strategy))] for _ in range(len(up_strategy))]

    for i in range(len(up_strategy)):
        for j in range(len(down_strategy)):
            var_up_strategy = up_strategy[i].out_strtg_group[idx_for_up]
            var_down_strategy = down_strategy[j].in_strtg_group[idx_for_down]
            comm_matrix[i][j] = calculate_resharding_cost(var, var_up_strategy, var_down_strategy,
                                                          device_mesh)

    return comm_matrix


def gen_mem_cost_matrix(var: MetaVar, up_strategy, down_strategy, device_mesh):
    comm_matrix = [[0 for _ in range(len(down_strategy))] for _ in range(len(up_strategy))]

    for i in range(len(up_strategy)):
        for j in range(len(down_strategy)):
            var_up_strategy = up_strategy[i]
            var_down_strategy = down_strategy[j]
            comm_matrix[i][j] = calculate_memory_cost(var, var_up_strategy, var_down_strategy,
                                                      device_mesh)

    return comm_matrix


# Lansong(TODO) generate_mem_matrix is deprecated, use gen_mem_cost_matrix instead
def generate_mem_matrix(var: MetaVar, up_strategy, down_strategy, idx_for_up, idx_for_down,
                        device_mesh):
    comm_matrix = [[0 for _ in range(len(down_strategy))] for _ in range(len(up_strategy))]

    for i in range(len(up_strategy)):
        for j in range(len(down_strategy)):
            var_up_strategy = up_strategy[i].out_strtg_group[idx_for_up]
            var_down_strategy = down_strategy[j].in_strtg_group[idx_for_down]
            comm_matrix[i][j] = calculate_memory_cost(var, var_up_strategy, var_down_strategy,
                                                      device_mesh)

    return comm_matrix


class ClusterMipInfo:

    def __init__(self, cluster: MetaNodeCluster, cluster_strtg_pool: ClusterStrategyPool,
                 mip_var: list):
        self.cluster = cluster
        self.cluster_strtg_pool = cluster_strtg_pool
        self.mip_var = mip_var

    def get_node_strtg(self, nd_id: int, strtg_idx: int) -> NodeSPMDStrategy:
        return self.cluster_strtg_pool.get_node_strtg(nd_id, strtg_idx)

    def __str__(self) -> str:
        res = str(self.cluster)
        res += "\n" + str(self.cluster_strtg_pool)
        return res

    def __repr__(self) -> str:
        return self.__str__()


class ClusterEdgeMipInfo:

    def __init__(self, edge: MetaVar):
        self.edge = edge
        self.up_node = None
        self.down_nodes = []
        self.idx_for_up = None
        self.idx_for_down = []
        self.comm_matrix = []
        self.mem_matrix = []
        self.mip_var = []

    def __str__(self) -> str:
        res = str(self.edge)
        return res

    def __repr__(self) -> str:
        return self.__str__()


class AutoFlowSolver:

    def __init__(self,
                 device_mesh=None,
                 constraints=None,
                 memory_ratio=0.9,
                 total_memery=None) -> None:
        self.m = mip.Model("autoflow")
        self.nodes = {}
        self.edges = {}
        self.device_mesh = device_mesh
        self.constraints = constraints

        self.memory_ratio = memory_ratio
        self.total_memery = 32 * 1024 * 1024 * 1024 if total_memery is None else total_memery
        self.max_memory_constrain = self.memory_ratio * self.total_memery

        self.clusters = {}
        self.cluster_edges = {}

    def add_coarsen_graph(self, graph: MetaGraph) -> None:
        self.graph = graph

        for cluster in graph.node_clusters:
            self.add_cluster(cluster)

        if mdconfig.log_level <= logging.DEBUG:
            for mip_info in self.clusters.values():
                print(str(mip_info))

    def add_cluster(self, cluster: MetaNodeCluster) -> None:
        cluster_id = cluster.unique_id
        cluster_strtg_pool = cluster.get_strtg_pool()
        cluster_strtg_num = cluster_strtg_pool.get_strtg_num()
        if cluster_strtg_num > 0:
            mip_vars = [self.m.add_var(var_type=mip.BINARY) for _ in range(cluster_strtg_num)]
            mip_info = ClusterMipInfo(cluster, cluster_strtg_pool, mip_vars)
            self.clusters[cluster_id] = mip_info

            # add edges from user side
            for input_node, input_idx in cluster.args.descs:
                if input_node.op_name == "placeholder" or input_node.op_name == "get_attr":
                    # placeholder or get_attr for a stateful variable has an implicit input
                    if input_node in self.graph.state_io_map:
                        output_var = self.graph.state_io_map[input_node]

                        # output var is the implicit input var of placeholder and get_attr
                        self.add_cluster_edge(output_var, input_idx, down_node=input_node)
                else:
                    var = input_node.invars[input_idx]
                    self.add_cluster_edge(var, input_idx, down_node=input_node)

            assert cluster.output_node
            for out_idx, var in enumerate(cluster.output_node.outvars):
                if var:
                    self.add_cluster_edge(var, out_idx, up_node=cluster.output_node)

    # Lansong(TODO) add_graph is deprecated, use add_coarsen_graph instead
    def add_graph(self, graph: MetaGraph) -> None:
        self.graph = graph

        if self.constraints:
            for cons in self.constraints:
                if cons.name not in [i.name for i in self.graph.input_list]:
                    continue
                cons_node = MetaNode("cons_" + cons.name, [], [cons], None)
                outvars_strategy = VarSPMDStrategyGroup(VarSPMDStrategy(self.constraints[cons]))
                self.nodes[cons_node.unique_key()] = {
                    "node": cons_node,
                    "strategy": NodeSPMDStrategy(None, outvars_strategy),
                    "mip_var": [self.m.add_var(var_type=mip.BINARY)]
                }
                self.add_edge(cons, up_node=cons_node)

        self.liveness = graph.liveness()
        if mdconfig.liveness_only_input:
            self.liveness = self.liveness[:1]

        for op in graph.op_list:
            self.add_node(op)

    # Lansong(TODO) add_node is deprecated, use add_cluster instead
    def add_node(self, node: MetaNode) -> None:
        unique_key_ = node.unique_key()

        strategies = node.get_strtg_pool().strategies
        if len(strategies) > 0:
            self.nodes[unique_key_] = {
                "node": node,
                "strategy": strategies,
                "mip_var": [self.m.add_var(var_type=mip.BINARY) for _ in range(len(strategies))]
            }

            for var in node.invars:
                self.add_edge(var, down_node=node)

            for var in node.outvars:
                if var:
                    self.add_edge(var, up_node=node)

    def add_cluster_edge(self, edge: MetaVar, io_idx: int, up_node=None, down_node=None) -> None:
        unique_key_ = edge.name

        if unique_key_ not in self.cluster_edges:
            edge_mip_info = ClusterEdgeMipInfo(edge=edge)
            self.cluster_edges[unique_key_] = edge_mip_info

        if up_node is not None:
            self.cluster_edges[unique_key_].up_node = up_node
            self.cluster_edges[unique_key_].idx_for_up = io_idx

        if down_node is not None:
            self.cluster_edges[unique_key_].down_nodes.append(down_node)
            self.cluster_edges[unique_key_].idx_for_down.append(io_idx)


        def _add_cluster_edge(up_node, idx_for_up, down_node, idx_for_down):
            up_cluster_strtg_pool = self.clusters[up_node.cluster_id].cluster_strtg_pool
            up_out_strategy_list = up_cluster_strtg_pool.get_outvar_strtg_list(
                up_node.unique_id, idx_for_up)

            down_cluster_strtg_pool = self.clusters[down_node.cluster_id].cluster_strtg_pool
            down_in_strategy_list = down_cluster_strtg_pool.get_invar_strtg_list(
                down_node.unique_id, idx_for_down)

            self.cluster_edges[unique_key_].mip_var.append([[
                self.m.add_var(var_type=mip.BINARY) for _ in range(len(down_in_strategy_list))
            ] for _ in range(len(up_out_strategy_list))])

            # calculate ``comm_matrix`` for this edge
            idx_for_up = self.cluster_edges[unique_key_].idx_for_up
            idx_for_down = self.cluster_edges[unique_key_].idx_for_down[-1]

            self.cluster_edges[unique_key_].comm_matrix.append(
                gen_comm_cost_matrix(
                    self.cluster_edges[unique_key_].edge,
                    up_out_strategy_list,
                    down_in_strategy_list,
                    self.device_mesh,
                ))

            self.cluster_edges[unique_key_].mem_matrix.append(
                gen_mem_cost_matrix(
                    self.cluster_edges[unique_key_].edge,
                    up_out_strategy_list,
                    down_in_strategy_list,
                    self.device_mesh,
                ))

        # if adding an up_node forms an edge
        if up_node and self.cluster_edges[unique_key_].down_nodes:
            for down_node, idx_for_down in zip(self.cluster_edges[unique_key_].down_nodes, 
                                               self.cluster_edges[unique_key_].idx_for_down):
                _add_cluster_edge(up_node, io_idx, down_node, idx_for_down)
        # if adding a down_node forms an edge
        elif self.cluster_edges[unique_key_].up_node and down_node:
            _add_cluster_edge(self.cluster_edges[unique_key_].up_node,
                              self.cluster_edges[unique_key_].idx_for_up,
                              down_node, io_idx)


    # Lansong(TODO) add_edge is deprecated, use add_cluster_edge instead
    def add_edge(self, edge: MetaVar, up_node=None, down_node=None) -> None:
        unique_key_ = edge.name

        if unique_key_ not in self.edges:
            self.edges[unique_key_] = {
                "edge": edge,
                "up_node": None,
                "down_node": [],
                "idx_for_up": None,
                "idx_for_down": [],
                "comm_matrix": [],
                "mem_matrix": [],
                "mip_var": [],
            }
        if up_node is not None:
            self.edges[unique_key_]["up_node"] = up_node.unique_key()
            self.edges[unique_key_]["idx_for_up"] = get_idx_in_var_list(edge, up_node.outvars)

        if down_node is not None:
            self.edges[unique_key_]["down_node"].append(down_node.unique_key())
            self.edges[unique_key_]["idx_for_down"].append(
                get_idx_in_var_list(edge, down_node.invars))

            if self.edges[unique_key_]["up_node"] is not None:

                up_node_key = self.edges[unique_key_]["up_node"]
                up_strategy = self.nodes[up_node_key]["strategy"]

                down_node_key = self.edges[unique_key_]["down_node"][-1]
                down_strategy = self.nodes[down_node_key]["strategy"]

                self.edges[unique_key_]["mip_var"].append(
                    [[self.m.add_var(var_type=mip.BINARY) for _ in range(len(down_strategy))]
                     for _ in range(len(up_strategy))])

                # calculate ``comm_matrix`` for this edge
                idx_for_up = self.edges[unique_key_]["idx_for_up"]
                idx_for_down = self.edges[unique_key_]["idx_for_down"][-1]

                self.edges[unique_key_]["comm_matrix"].append(
                    generate_comm_matrix(
                        self.edges[unique_key_]["edge"],
                        up_strategy,  # list of NodeSPMDStrategy
                        down_strategy,  # list of NodeSPMDStrategy
                        idx_for_up,
                        idx_for_down,
                        self.device_mesh,
                    ))

                self.edges[unique_key_]["mem_matrix"].append(
                    generate_mem_matrix(
                        self.edges[unique_key_]["edge"],
                        up_strategy,  # list of NodeSPMDStrategy
                        down_strategy,  # list of NodeSPMDStrategy
                        idx_for_up,
                        idx_for_down,
                        self.device_mesh,
                    ))

    # Lansong(TODO) ilp_optimize is deprecated, use ilp_solve instead
    def ilp_optimize(self, count_invars=False):
        comm_cost, mem_cost = 0, 0
        for edge in self.edges.values():
            for idx in range(len(edge["mip_var"])):
                mip_var = edge["mip_var"][idx]
                comm_matrix = edge["comm_matrix"][idx]
                mem_matrix = edge["mem_matrix"][idx]
                shape_1 = len(mip_var)
                shape_2 = len(mip_var[0])
                comm_cost = comm_cost + mip.xsum(mip_var[i][j] * comm_matrix[i][j]
                                                 for i in range(shape_1) for j in range(shape_2))
                mem_cost = mem_cost + mip.xsum(mip_var[i][j] * mem_matrix[i][j]
                                               for i in range(shape_1) for j in range(shape_2))

            def _mem_cost(var_size, down_strategy, idx_for_down):
                memory_cost_list = []
                for i in range(len(down_strategy)):
                    strategy = down_strategy[i].in_strtg_group[idx_for_down]

                    # FIXME: only use out_strategy here for shard_size is ok?
                    s1, s2 = strategy[0], strategy[1]
                    shard_size = 1
                    if s1.state == SPMD.SHARD:
                        shard_size *= self.device_mesh[0]
                    if s2.state == SPMD.SHARD:
                        shard_size *= self.device_mesh[1]

                    memory_cost_list.append(var_size // shard_size)

                return memory_cost_list

            if count_invars and edge["up_node"] is None:
                var_size = edge["edge"].get_var_size()
                for down_node_key, idx_for_down in zip(edge["down_node"], edge["idx_for_down"]):
                    down_strategy = self.nodes[down_node_key]["strategy"]
                    __mem_cost = _mem_cost(var_size, down_strategy, idx_for_down)
                    down_node_mip_var = self.nodes[down_node_key]["mip_var"]
                    mem_cost = mem_cost + mip.xsum(down_node_mip_var[i] * __mem_cost[i]
                                                   for i in range(len(down_node_mip_var)))

        for edge in self.edges.values():
            for idx in range(len(edge["mip_var"])):
                mip_var = edge["mip_var"][idx]
                shape_1 = len(mip_var)
                shape_2 = len(mip_var[0])
                self.m += mip.xsum(mip_var[i][j] for i in range(shape_1)
                                   for j in range(shape_2)) == 1

                up_node_key = edge["up_node"]
                up_node_mip_var = self.nodes[up_node_key]["mip_var"]

                down_node_key = edge["down_node"][idx]
                down_node_mip_var = self.nodes[down_node_key]["mip_var"]

                for i in range(shape_1):
                    for j in range(shape_2):
                        self.m += mip_var[i][j] <= up_node_mip_var[i]
                        self.m += mip_var[i][j] <= down_node_mip_var[j]
                        self.m += mip_var[i][j] >= up_node_mip_var[i] + down_node_mip_var[j] - 1

        idx_record = {}
        for idx_ in range(len(self.liveness)):
            mem_live = self.liveness[idx_]
            op_ = self.graph.op_list[idx_]

            mip_var_list = []
            mem_matrix_list = []
            for tensor_name in mem_live:
                if tensor_name not in self.edges:
                    continue
                if tensor_name not in idx_record:
                    idx_record[tensor_name] = 0
                if len(self.edges[tensor_name]["mip_var"]) > 0:
                    edge_idx = min(
                        len(self.edges[tensor_name]["mip_var"]) - 1, idx_record[tensor_name])
                    mip_var_list.append(self.edges[tensor_name]["mip_var"][edge_idx])
                    mem_matrix_list.append(self.edges[tensor_name]["mem_matrix"][edge_idx])
                if self.edges[tensor_name]["up_node"] is None:
                    down_node_key, idx_for_down = self.edges[tensor_name]["down_node"][
                        0], self.edges[tensor_name]["idx_for_down"][0]
                    var_size = self.edges[tensor_name]["edge"].get_var_size()
                    down_strategy = self.nodes[down_node_key]["strategy"]
                    __mem_cost = _mem_cost(var_size, down_strategy, idx_for_down)
                    mip_var_list.append([self.nodes[down_node_key]["mip_var"]])
                    mem_matrix_list.append([__mem_cost])

            for var in op_.invars:
                if var.name not in self.edges:
                    continue
                if var.name not in idx_record:
                    idx_record[var.name] = 0
                idx_record[var.name] += 1

            if len(mip_var_list) >= 1:
                need_sum = []
                for mip_var, mem_matrix in zip(mip_var_list, mem_matrix_list):
                    shape_1 = len(mip_var)
                    shape_2 = len(mip_var[0])
                    for i in range(shape_1):
                        for j in range(shape_2):
                            need_sum.append(mip_var[i][j] * mem_matrix[i][j])
                self.m += mip.xsum(i for i in need_sum) <= self.max_memory_constrain

        for node in self.nodes.values():
            mip_var = node["mip_var"]
            shape_1 = len(mip_var)
            self.m += mip.xsum(mip_var[i] for i in range(shape_1)) == 1

        self.m.objective = mip.minimize(comm_cost + 0.00000001 * mem_cost)

        self.m.verbose = 0
        status = self.m.optimize(max_seconds_same_incumbent=mdconfig.max_seconds_same_incumbent)
        logger.info(f'[AutoFlowSolver.status]:\t {status}')
        logger.info(f'[AutoFlowSolver.solution_cost]:\t {self.m.objective_value}')

        return self.get_optimal_stratey()

    # Lansong(TODO) get_optimal_stratey is deprecated
    def get_optimal_stratey(self):
        optimal_stratey = {}
        for unique_key_ in self.nodes:
            node = self.nodes[unique_key_]['node']
            opt_ = None
            strategy_list = self.nodes[unique_key_]['strategy']
            mip_var = self.nodes[unique_key_]['mip_var']
            assert len(strategy_list) == len(mip_var)
            for s_, mip_var_s in zip(strategy_list, mip_var):
                if mip_var_s.x == 1:
                    opt_ = s_
            optimal_stratey[unique_key_] = {'node': node, 'strategy': opt_}

        optimal_stratey = self.strategies_compact_output(optimal_stratey)

        if mdconfig.log_level <= logging.DEBUG:
            print("optimal_stratey:")
            pprint(optimal_stratey)

        return optimal_stratey

    def ilp_solve(self, count_invars=False):
        comm_cost, mem_cost = 0, 0
        for edge_mip_info in self.cluster_edges.values():
            for idx in range(len(edge_mip_info.mip_var)):
                mip_var = edge_mip_info.mip_var[idx]
                comm_matrix = edge_mip_info.comm_matrix[idx]
                mem_matrix = edge_mip_info.mem_matrix[idx]
                shape_1 = len(mip_var)
                shape_2 = len(mip_var[0])
                comm_cost = comm_cost + mip.xsum(mip_var[i][j] * comm_matrix[i][j]
                                                 for i in range(shape_1) for j in range(shape_2))
                mem_cost = mem_cost + mip.xsum(mip_var[i][j] * mem_matrix[i][j]
                                               for i in range(shape_1) for j in range(shape_2))

            def _mem_cost(var_size, down_strategy: 'list[VarSPMDStrategy]'):
                memory_cost_list = []
                for i in range(len(down_strategy)):
                    strategy = down_strategy[i]

                    # FIXME: only use out_strategy here for shard_size is ok?
                    s1, s2 = strategy[0], strategy[1]
                    shard_size = 1
                    if s1.state == SPMD.SHARD:
                        shard_size *= self.device_mesh[0]
                    if s2.state == SPMD.SHARD:
                        shard_size *= self.device_mesh[1]

                    memory_cost_list.append(var_size // shard_size)

                return memory_cost_list

            if count_invars and edge_mip_info.up_node is None:
                var_size = edge_mip_info.edge.get_var_size()
                for down_node, idx_for_down in zip(edge_mip_info.down_nodes,
                                                   edge_mip_info.idx_for_down):
                    down_cluster_mip_info = self.clusters[down_node.cluster_id]
                    down_cluster_strtg_pool = down_cluster_mip_info.cluster_mip_info
                    down_strategy = down_cluster_strtg_pool.get_invar_strtg_list(
                        down_node.unique_id, idx_for_down)
                    __mem_cost = _mem_cost(var_size, down_strategy)
                    down_cluster_mip_var = self.clusters[down_node.cluster_id].mip_var
                    mem_cost = mem_cost + mip.xsum(down_cluster_mip_var[i] * __mem_cost[i]
                                                   for i in range(len(down_cluster_mip_var)))

        for edge_mip_info in self.cluster_edges.values():
            for idx in range(len(edge_mip_info.mip_var)):
                mip_var = edge_mip_info.mip_var[idx]
                shape_1 = len(mip_var)
                shape_2 = len(mip_var[0])
                self.m += mip.xsum(mip_var[i][j] for i in range(shape_1)
                                   for j in range(shape_2)) == 1

                up_cluster_id = edge_mip_info.up_node.cluster_id
                up_cluster_mip_var = self.clusters[up_cluster_id].mip_var

                down_node = edge_mip_info.down_nodes[idx]
                down_cluster_mip_var = self.clusters[down_node.cluster_id].mip_var

                for i in range(shape_1):
                    for j in range(shape_2):
                        self.m += mip_var[i][j] <= up_cluster_mip_var[i]
                        self.m += mip_var[i][j] <= down_cluster_mip_var[j]
                        self.m += mip_var[i][
                            j] >= up_cluster_mip_var[i] + down_cluster_mip_var[j] - 1


#        Lansong(TODO) support memory limit
#        idx_record = {}
#        for idx_ in range(len(self.liveness)):
#            mem_live = self.liveness[idx_]
#            op_ = self.graph.op_list[idx_]
#
#            mip_var_list = []
#            mem_matrix_list = []
#            for tensor_name in mem_live:
#                if tensor_name not in self.cluster_edges:
#                    continue
#                if tensor_name not in idx_record:
#                    idx_record[tensor_name] = 0
#                if len(self.cluster_edges[tensor_name]["mip_var"]) > 0:
#                    edge_idx = min(
#                        len(self.cluster_edges[tensor_name]["mip_var"]) - 1, idx_record[tensor_name])
#                    mip_var_list.append(self.cluster_edges[tensor_name]["mip_var"][edge_idx])
#                    mem_matrix_list.append(self.cluster_edges[tensor_name]["mem_matrix"][edge_idx])
#                if self.cluster_edges[tensor_name]["up_node"] is None:
#                    down_node_key, idx_for_down = self.cluster_edges[tensor_name]["down_nodes"][
#                        0], self.cluster_edges[tensor_name]["idx_for_down"][0]
#                    var_size = self.cluster_edges[tensor_name]["edge"].get_var_size()
#                    down_strategy = self.clusters[down_node_key]["strategy"]
#                    __mem_cost = _mem_cost(var_size, down_strategy, idx_for_down)
#                    mip_var_list.append([self.clusters[down_node_key]["mip_var"]])
#                    mem_matrix_list.append([__mem_cost])
#
#            for var in op_.invars:
#                if var.name not in self.cluster_edges:
#                    continue
#                if var.name not in idx_record:
#                    idx_record[var.name] = 0
#                idx_record[var.name] += 1
#
#            if len(mip_var_list) >= 1:
#                need_sum = []
#                for mip_var, mem_matrix in zip(mip_var_list, mem_matrix_list):
#                    shape_1 = len(mip_var)
#                    shape_2 = len(mip_var[0])
#                    for i in range(shape_1):
#                        for j in range(shape_2):
#                            need_sum.append(mip_var[i][j] * mem_matrix[i][j])
#                self.m += mip.xsum(i for i in need_sum) <= self.max_memory_constrain

        for cluster_mip_info in self.clusters.values():
            mip_var = cluster_mip_info.mip_var
            shape_1 = len(mip_var)
            self.m += mip.xsum(
                mip_var[i]
                for i in range(shape_1)) == 1  # one and only one cluster strategy is active

        self.m.objective = mip.minimize(comm_cost + 0.00000001 * mem_cost)

        self.m.verbose = 0
        status = self.m.optimize(max_seconds_same_incumbent=mdconfig.max_seconds_same_incumbent)
        logger.info(f'[AutoFlowSolver.status]:\t {status}')
        logger.info(f'[AutoFlowSolver.solution_cost]:\t {self.m.objective_value}')

        graph_strategy = self.get_strategies()
        total_comm_cost = self.calc_graph_comm_cost(graph_strategy)
        logger.info(f'[Communication Cost]:\t {total_comm_cost}')

        return self.strategies_compact_output(graph_strategy)

    def strategies_compact_output(self, optimal_strategies):

        for node_name in optimal_strategies:
            nstrtgy = optimal_strategies[node_name]
            out_strategies_compact = []
            for idx in nstrtgy['node'].compact_out_idx_tbl:
                if idx < 0:
                    out_strategies_compact.append(None)
                else:
                    out_strategies_compact.append(nstrtgy['strategy'].out_strtg_group[idx])
            
            nstrtgy['strategy'].out_strtg_group = VarSPMDStrategyGroup(*out_strategies_compact)

        return optimal_strategies

    def get_strategies(self):
        optimal_strategies = {}

        for cluster_mip_info in self.clusters.values():
            cluster = cluster_mip_info.cluster
            opt_strtg_idx = -1
            for strtg_idx, mip_var_s in enumerate(cluster_mip_info.mip_var):
                if mip_var_s.x == 1:
                    opt_strtg_idx = strtg_idx
                    break
            assert opt_strtg_idx >= 0

            for node in cluster.nodes.values():
                nd_strtg = cluster_mip_info.get_node_strtg(node.unique_id, opt_strtg_idx)
                optimal_strategies[node.unique_key()] = {'node': node, 'strategy': nd_strtg}

        if mdconfig.log_level <= logging.DEBUG:
            print("optimal_strategies:")
            pprint(optimal_strategies)

        return optimal_strategies

    def calc_graph_comm_cost(self, graph_strategy):
        total_comm_cost = 0
        default_strtg = [SPMD(SPMD.REPLICATE) for _ in range(len(self.device_mesh))]
        for node in self.graph.op_list:
            for in_idx, in_var in enumerate(node.invars):
                if in_var.up_node.unique_key() in graph_strategy:
                    up_node_strtg = graph_strategy[in_var.up_node.unique_key()]['strategy']
                    var_up_strtg = up_node_strtg.get_outvar_strtg(in_var.idx_for_up)
                else:
                    var_up_strtg = VarSPMDStrategy(*default_strtg)

                if node.unique_key() in graph_strategy:
                    node_strtg = graph_strategy[node.unique_key()]['strategy']
                    var_down_strtg = node_strtg.get_invar_strtg(in_idx)
                else:
                    var_down_strtg = VarSPMDStrategy(*default_strtg)

                comm_cost = calculate_resharding_cost(in_var, var_up_strtg, var_down_strtg,
                                                      self.device_mesh)

                total_comm_cost += comm_cost
                if mdconfig.log_level <= logging.DEBUG:
                    if comm_cost > 0:
                        print(("cost between {} and {} is {}").format(in_var.up_node.name, node.name, comm_cost))

        return total_comm_cost

    def beam_search(self, candidate_num=100):

        def get_new_cost(strategy, node, strategy_idx):
            cost = 0.
            edge_list = [self.edges[invar.name] for invar in node['node'].invars]
            for edge in edge_list:
                up_node_key = edge["up_node"]

                idx = edge["down_node"].index(node['node'].unique_key())
                if up_node_key in strategy:
                    up_node_strategy_idx = strategy[up_node_key]['strategy_idx']
                    mem_cost = edge["mem_matrix"][idx][up_node_strategy_idx][strategy_idx]
                    comm_cost = edge["comm_matrix"][idx][up_node_strategy_idx][strategy_idx]
                    cost += comm_cost

            return cost

        def add_candidate(strategy_candidate, accumulate_cost, node):
            new_strategy_candidate = []
            new_accumulate_cost = []
            key_ = node['node'].unique_key()
            if len(strategy_candidate) == 0:
                for idx in range(len(node['strategy'])):
                    stratey = {key_: {'node': node['node'], 'strategy_idx': idx}}
                    new_strategy_candidate.append(stratey)
                    new_accumulate_cost.append(0.)
            else:
                for idx, strategy in enumerate(strategy_candidate):
                    old_cost = accumulate_cost[idx]
                    for idx in range(len(node['strategy'])):
                        new_cost = get_new_cost(strategy, node, idx)
                        new_strategy = {key: strategy[key] for key in strategy}
                        new_strategy[key_] = {'node': node['node'], 'strategy_idx': idx}
                        new_strategy_candidate.append(new_strategy)
                        new_accumulate_cost.append(old_cost + new_cost)

            return new_strategy_candidate, new_accumulate_cost

        def select_candidate(strategy_candidate, accumulate_cost, candidate_num):
            assert len(strategy_candidate) == len(accumulate_cost)
            if len(accumulate_cost) <= candidate_num:
                return strategy_candidate, accumulate_cost

            accumulate_cost, strategy_candidate = shuffle_list(accumulate_cost, strategy_candidate)
            accumulate_cost, strategy_candidate = zip(
                *sorted(zip(accumulate_cost, strategy_candidate), key=lambda x: x[0]))

            return strategy_candidate[:candidate_num], accumulate_cost[:candidate_num]

        strategy_candidate = []
        accumulate_cost = []
        for unique_key_ in self.nodes:
            node = self.nodes[unique_key_]

            strategy_candidate, accumulate_cost = add_candidate(strategy_candidate,
                                                                accumulate_cost, node)

            strategy_candidate, accumulate_cost = select_candidate(strategy_candidate,
                                                                   accumulate_cost, candidate_num)

            accumulate_cost, strategy_candidate = zip(
                *sorted(zip(accumulate_cost, strategy_candidate), key=lambda x: x[0]))

        strategy = strategy_candidate[0]

        optimal_stratey = {}
        for key in strategy:
            node = strategy[key]['node']
            node_strategy_list = self.nodes[node.unique_key()]['strategy']
            optimal_stratey[key] = {
                "node": node,
                "strategy": node_strategy_list[strategy[key]['strategy_idx']]
            }

        logger.info(f'=========== solution cost:\t {accumulate_cost[0]}')

        return optimal_stratey

easydist/autoflow/solver.py
alibaba/easydist
# Copyright (c) 2023, Alibaba Group;
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================

import os
import logging

def _get_env_or_raise(env_var: str) -> str:
    env_val = os.environ.get(env_var, None)
    if not env_val:
        raise ValueError("environment variable {env_var} expected, but not set")
    else:
        return env_val

log_level_str = os.environ.get("EASYDIST_LOGLEVEL", 'INFO').upper()
log_level = logging.getLevelName(log_level_str)

dump_dir = os.environ.get("EASYDIST_DUMP_PATH", './md_dump')
enable_compile_cache = os.environ.get("ENABLE_COMPILE_CACHE", "False").upper() in ["1", "TRUE"]
compile_cache_dir = os.environ.get("EASYDIST_COMPILE_CACHE_PATH", './md_compiled')

easydist_device = os.environ.get("EASYDIST_DEVICE", "cuda")

# MetaSPMD Annotation

use_hint = os.environ.get("EASYDIST_USE_HINT") == "1"
extend_space = False

# Solver

liveness_only_input = False
max_seconds_same_incumbent = float('inf')

all_to_all_punish_factor = 3.

enable_graph_coarsen = os.environ.get("ENABLE_GRAPH_COARSEN", "True").upper() in ["1", "TRUE"]
coarsen_level = int(os.environ.get("COARSEN_LEVEL", "1"))

# Master address and port
master_addr = os.environ.get("MASTER_ADDR", "localhost")
master_port = int(_get_env_or_raise("MASTER_PORT"))

# PyTorch

comm_optimization = False
override_dtensor_rule = False

# runtime

use_contiguous_buffer = False

easydist/config.py
alibaba/easydist
# Copyright (c) 2023, Alibaba Group;
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================

import os
import logging

import jax
from mpi4py import MPI

from .api import easydist_shard, get_opt_strategy, set_device_mesh
from .sharding_interpreter import EDJaxShardingAnn
from .bridge import jax2md_bridge

__all__ = [
    "easydist_shard", "get_opt_strategy", "set_device_mesh", "EDJaxShardingAnn", "jax2md_bridge"
]

logger = logging.getLogger(__name__)


def is_jax_distributed_initialized():
    return jax._src.distributed.global_state.client is not None


def easydist_setup_jax(device, allow_tf32):
    os.environ["NVIDIA_TF32_OVERRIDE"] = "1" if allow_tf32 else "0"
    jax.config.update('jax_platforms', device)

    # setup distributed
    comm = MPI.COMM_WORLD
    size, rank = comm.Get_size(), comm.Get_rank()

    if not is_jax_distributed_initialized():

        jax.distributed.initialize(coordinator_address="localhost:19705",
                                   num_processes=size,
                                   process_id=rank,
                                   local_device_ids=rank)

        logging.info(
            f"[Rank {rank}], Global Devices: {jax.device_count()}, Local Devices: {jax.local_device_count()}"
        )

easydist/jax/__init__.py
alibaba/easydist
# Copyright (c) 2023, Alibaba Group;
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================

import functools
import logging
import time
from typing import Any

import jax
import rich
from jax import core
from jax._src.distributed import global_state
from jax._src.util import safe_map
from jax.experimental import mesh_utils, multihost_utils
from jax.sharding import Mesh, NamedSharding, PartitionSpec

import easydist.config as mdconfig
from easydist.autoflow import AutoFlowSolver
from easydist.metashard import metair

from .device_mesh import set_device_mesh, get_device_mesh
from .bridge import jax2md_bridge
from .sharding_interpreter import EDJaxShardingAnn
from .utils import _sharding_ann_env

logger = logging.getLogger(__name__)


INPUT_STRATEGY = None


def to_shape_array(x):
    if isinstance(x, jax.Array) and not jax.core.is_opaque_dtype(x.dtype):
        return core.ShapedArray(shape=x.shape, dtype=x.dtype)
    else:
        return x


def materialize(x):
    if isinstance(x, core.ShapedArray):
        key = jax.random.PRNGKey(seed=42)
        if x.dtype.name in ["float64", "float32", "float16"]:
            return jax.random.normal(key, shape=x.shape, dtype=x.dtype)
        elif x.dtype.name in ["int32", "unint32", "int64", "uint64", "uint8"]:
            return jax.random.randint(key, shape=x.shape, dtype=x.dtype, minval=1, maxval=8)
        elif x.dtype.name in ["bool"]:
            return jax.random.normal(key, shape=x.shape) > 1.
        else:
            return jax.numpy.empty(shape=x.shape, dtype=x.dtype)
    return x


def convert(strategy, mesh, val):
    axis_names = mesh.axis_names

    s1, s2 = strategy[0], strategy[1]
    ndim = len(val.shape)
    mesh_shape = [None] * ndim

    # we use strategy except val.shape -> ()
    if ndim > 0:
        for idx, s_ in enumerate([s1, s2]):
            if s_.state == metair.SPMD.SHARD:
                dim = s_.args["dim"]
                if mesh_shape[dim] == None:
                    mesh_shape[dim] = axis_names[idx]
                else:
                    mesh_shape[dim] = mesh.axis_names

    return NamedSharding(mesh, PartitionSpec(*mesh_shape))


def _get_local_array(array, specs):
    mesh = get_device_mesh()
    mesh_shape = mesh.device_ids.shape
    axis_names = mesh.axis_names
    for idx in range(len(specs)):
        for axis_idx, axis in enumerate(axis_names):
            if specs[idx] == axis:
                global_rank = global_state.process_id
                rank_coords = (mesh.device_ids == global_rank).nonzero()
                mesh_dim_rank = rank_coords[axis_idx].item()
                array = jax.numpy.array_split(array, mesh_shape[axis_idx], axis=idx)[mesh_dim_rank]

    return array


def shard_module(flatten_args):

    device_mesh = get_device_mesh()

    for i in range(len(flatten_args)):
        if isinstance(flatten_args[i], jax.Array) or isinstance(flatten_args[i], jax.ShapedArray):
            strategy = convert(INPUT_STRATEGY[i], device_mesh, flatten_args[i])
            flatten_args[i] = materialize(flatten_args[i])
            local_array = _get_local_array(flatten_args[i], strategy.spec)
            flatten_args[i] = multihost_utils.host_local_array_to_global_array(
                local_array, device_mesh, strategy.spec)

    return flatten_args


def add_sharding_jaxpr(jaxpr, consts, shard_strategy, args):
    env = {}

    def read(var):
        if type(var) is core.Literal:
            return var.val
        return env[var]

    def write(var, val):
        env[var] = val

    # Args now correspond to Jaxpr outvars
    safe_map(write, jaxpr.invars, args)
    safe_map(write, jaxpr.constvars, consts)

    # Looping backward
    for eqn in jaxpr.eqns:
        subfuns, bind_params = eqn.primitive.get_bind_params(eqn.params)

        onelevel_key = eqn.primitive.__str__()
        if eqn.primitive.__str__() == "custom_jvp_call":
            onelevel_key += "[" + \
                subfuns[0].f.args[0].eqns[0].params['name'] + "]"
        if eqn.primitive.__str__() == "xla_call":
            onelevel_key += "[" + eqn.params['name'] + "]"

        invars = [var for var in eqn.invars if isinstance(var, jax.core.Var)]

        #  outvars are now invars
        invals = safe_map(read, eqn.invars)

        # NOTE, use last outvars for name of MetaNode in Jax
        unique_key = str(eqn.outvars[-1])

        mesh = get_device_mesh()

        if unique_key in shard_strategy:
            in_strtg_group = shard_strategy[unique_key]['strategy'].in_strtg_group

            sharding_idx = 0
            for idx in range(len(invals)):
                if isinstance(invals[idx], jax.interpreters.partial_eval.DynamicJaxprTracer):
                    strategy_ = convert(in_strtg_group.get_var_strtg(sharding_idx), mesh,
                                        invals[idx])
                    invals[idx] = jax.lax.with_sharding_constraint(invals[idx], strategy_)
                    sharding_idx += 1

        outval = eqn.primitive.bind(*subfuns, *invals, **bind_params)

        if isinstance(outval, jax.interpreters.partial_eval.DynamicJaxprTracer):
            outval = [outval]

        safe_map(write, eqn.outvars, outval)

    outvals = safe_map(read, jaxpr.outvars)

    return outvals


def get_opt_strategy(func, *args, **kwargs):

    global INPUT_STRATEGY

    closed_jaxpr = jax.make_jaxpr(func)(*args, **kwargs)

    # use tf32 during MetaSPMD Annotaion
    with _sharding_ann_env():

        start_t = time.perf_counter()
        sharding_interpreter = EDJaxShardingAnn(closed_jaxpr.jaxpr)
        sharding_info, shape_info = sharding_interpreter.run(closed_jaxpr.literals, *args,
                                                             **kwargs)
        logger.info(f"[EDJaxShardingAnn.run]: {time.perf_counter() - start_t} s.")

        if mdconfig.log_level <= logging.DEBUG and global_state.process_id == 0:
            rich.print("sharding_info:\n", sharding_info)
            rich.print("shape_info:\n", shape_info)

    meta_graph = jax2md_bridge(closed_jaxpr.jaxpr, sharding_info, shape_info)

    if mdconfig.log_level <= logging.DEBUG and global_state.process_id == 0:
        rich.print(meta_graph)

    device_mesh = get_device_mesh()
    solver = AutoFlowSolver(device_mesh.device_ids.shape)

    if mdconfig.enable_graph_coarsen:
        logger.info(f"enable graph coarsen with level {mdconfig.coarsen_level}.")
        solver.add_coarsen_graph(meta_graph)
    else:
        solver.add_graph(meta_graph)

    start_t = time.perf_counter()
    if mdconfig.enable_graph_coarsen:
        opt_strategy = solver.ilp_solve()
    else:
        opt_strategy = solver.ilp_optimize()

    logger.info(f"[AutoFlowSolver.ilp_optimize]: {time.perf_counter() - start_t} s.")
    # start_t = time.perf_counter()
    # beam_search_strategy = solver.beam_search()
    # logger.info(f"[AutoFlowSolver.beam_search]: {time.perf_counter() - start_t} s.")

    INPUT_STRATEGY = meta_graph.get_input_strategy(opt_strategy)

    if mdconfig.log_level <= logging.DEBUG and global_state.process_id == 0:
        rich.print(opt_strategy)

    return opt_strategy


def easydist_shard(fun, shard_strategy={}):

    @functools.wraps(fun)
    def wrapped(*args, **kwargs):
        # Since we assume unary functions, we won't worry about flattening and
        # unflattening arguments.
        closed_jaxpr, return_shape = jax.make_jaxpr(fun, return_shape=True)(*args, **kwargs)

        flatten_args, in_tree = jax.tree_util.tree_flatten(args)

        out = add_sharding_jaxpr(closed_jaxpr.jaxpr, closed_jaxpr.literals, shard_strategy,
                                 flatten_args)

        flatten_shape, shape_tree = jax.tree_util.tree_flatten(return_shape)

        out_idx = 0
        for idx, shape in enumerate(flatten_shape):
            if isinstance(shape, jax.ShapeDtypeStruct):
                flatten_shape[idx] = out[out_idx]
                out_idx += 1

        out_tree = jax.tree_util.tree_unflatten(shape_tree, flatten_shape)

        return out_tree

    return wrapped


# =============== Experimetal API for Jax ==================


def _compile(func, args, kwargs):

    # setup the device mesh
    size = jax.device_count()
    devices = mesh_utils.create_device_mesh((1, size))
    mesh = Mesh(devices, axis_names=('a', 'b'))

    set_device_mesh(mesh)

    opt_strategy = get_opt_strategy(func, *args, **kwargs)

    shard_func = jax.jit(easydist_shard(func, opt_strategy))

    if mdconfig.log_level <= logging.DEBUG and global_state.process_id == 0:
        closed_jaxpr = jax.make_jaxpr(shard_func)(*args, **kwargs)
        print(closed_jaxpr.jaxpr)
        print(closed_jaxpr.literals)

    def compiled_func_return(*args, **kwargs):

        flatten_args, specs = jax.tree_util.tree_flatten((args, kwargs))
        flatten_args = shard_module(flatten_args)
        shard_args, shard_kwargs = jax.tree_util.tree_unflatten(specs, flatten_args)

        with jax.spmd_mode('allow_all'):
            out = shard_func(*shard_args, **shard_kwargs)

        return out

    return compiled_func_return


class CompiledFuncWrapper:

    def __init__(self,
                 func,
                 compile_only=False) -> None:
        self.original_func = func
        self.compiled_func = None
        self.compile_only = compile_only

    def __call__(self, *args: Any, **kwargs: Any) -> Any:
        if self.compiled_func is None:
            self.compiled_func = _compile(self.original_func, args, kwargs)

        if self.compile_only:
            return self.compiled_func

        return self.compiled_func(*args, **kwargs)


def easydist_compile(func=None,
                     max_solver_time=float("inf"),
                     liveness_only_input=False,
                     compile_only=False):

    mdconfig.liveness_only_input = liveness_only_input
    mdconfig.max_seconds_same_incumbent = max_solver_time

    if func:
        return CompiledFuncWrapper(func, compile_only)

    else:

        def wrapper(func):
            return CompiledFuncWrapper(func, compile_only)

        return wrapper

easydist/jax/api.py
alibaba/easydist
# Copyright (c) 2023, Alibaba Group;
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================

import jax

from easydist.metashard.metair import MetaGraph, MetaNode, MetaVar
import easydist.config as mdconfig


def jax2md_bridge(jaxpr: jax.core.Jaxpr, sharding_info, meta_info) -> MetaGraph:

    meta_graph = MetaGraph(jaxpr)

    meta_var_map = {}

    for in_var in jaxpr.invars + jaxpr.constvars:
        meta_var = MetaVar(name=in_var.__str__(),
                           shape=meta_info[in_var.__str__()]["shape"],
                           dtype=meta_info[in_var.__str__()]["dtype"])
        meta_var_map[in_var.__str__()] = meta_var

        node_sharding_info = None
        onelevel_key = "placeholder"
        twolevel_key = str([(meta_var.shape, meta_var.dtype)])
        if onelevel_key in sharding_info:
            node_sharding_info = sharding_info[onelevel_key].get(twolevel_key, None)

        meta_node = MetaNode(name=in_var.__str__(),
                             op_name="placeholder",
                             invars=[],
                             outvars=[meta_var],
                             sharding_info=node_sharding_info,
                             is_placeholder=True)
        meta_node.compact_out_idx_tbl = [0]
        meta_node.compact_in_idx_tbl = []

        meta_graph.add_node(meta_node)
        meta_graph.add_input(meta_node)

    for eqn in jaxpr.eqns:
        subfuns, bind_params = eqn.primitive.get_bind_params(eqn.params)
        onelevel_key = eqn.primitive.__str__()
        if eqn.primitive.__str__() == "custom_jvp_call":
            onelevel_key += "[" + subfuns[0].f.args[0].eqns[0].params['name'] + "]"
        if eqn.primitive.__str__() in ["xla_call", "pjit"]:
            onelevel_key += "[" + eqn.params['name'] + "]"

        node_sharding_info = None
        if onelevel_key in sharding_info:
            abstract_list = []
            for var in eqn.invars:
                if var.__str__() in meta_info:
                    abstract_list.append(
                        (meta_info[var.__str__()]["shape"], meta_info[var.__str__()]["dtype"]))
                elif type(var) is jax.core.Literal:
                    abstract_list.append(var.val)
                else:
                    abstract_list.append(var)
            twolevel_key = str(abstract_list) + str(bind_params)
            if twolevel_key in sharding_info[onelevel_key]:
                node_sharding_info = sharding_info[onelevel_key][twolevel_key]

        outvars_ = []
        for var in eqn.outvars:
            name = var.__str__()
            outvar = MetaVar(name=name,
                             shape=meta_info[name]["shape"],
                             dtype=meta_info[name]["dtype"])
            outvars_.append(outvar)
            meta_var_map[name] = outvar

        # NOTE, use last outvars for name of MetaNode in Jax
        meta_node = MetaNode(name=name,
                             op_name=onelevel_key,
                             invars=[
                                 meta_var_map[var.__str__()] for var in eqn.invars
                                 if isinstance(var, jax.core.Var)
                             ],
                             outvars=outvars_,
                             sharding_info=node_sharding_info)

        meta_node.compact_out_idx_tbl = list(range(len(outvars_)))

        compact_in_idx_tbl = []
        compact_in_idx = 0
        for var in eqn.invars:
            if isinstance(var, jax.core.Var):
                compact_in_idx_tbl.append(compact_in_idx)
                compact_in_idx = compact_in_idx + 1
            else:
                compact_in_idx_tbl.append(-1)
        meta_node.compact_in_idx_tbl = compact_in_idx_tbl
        meta_graph.add_node(meta_node)

    for var in jaxpr.outvars:
        meta_graph.add_output(meta_var_map[var.__str__()])

    meta_graph.coarsen(coarsen_level=mdconfig.coarsen_level)

    return meta_graph

easydist/jax/bridge.py
alibaba/easydist
# Copyright (c) 2023, Alibaba Group;
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================

import functools
import operator

from easydist.metashard import metair

JAX_DEVICE_MESH = None

def set_device_mesh(device_mesh):
    global JAX_DEVICE_MESH
    JAX_DEVICE_MESH = device_mesh

    mesh_shape = device_mesh.device_ids.shape
    if mesh_shape[0] == 1:
        metair.DEVICE_MESH_1D = 0
    elif mesh_shape[1] == 1:
        metair.DEVICE_MESH_1D = 1


def get_device_mesh():
    global JAX_DEVICE_MESH
    return JAX_DEVICE_MESH


def device_mesh_world_size(device_mesh=None):
    if device_mesh is None:
        device_mesh = get_device_mesh()

    if device_mesh is None:
        return None

    device_mesh_shape = device_mesh.device_ids.shape

    return functools.reduce(operator.mul, device_mesh_shape)

easydist/jax/device_mesh.py
alibaba/easydist
# Copyright (c) 2023, Alibaba Group;
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================

import re
from typing import Any

import jax
from jax._src import source_info_util
from jax._src.util import safe_map
from jax._src.core import ShapedArray

from easydist.metashard import MetaOp, view_propagation
from .device_mesh import device_mesh_world_size


def generate_str(invals, bind_params):
    return str([(val.shape, val.dtype.name) if isinstance(val, jax.Array) else val
                for val in invals]) + str(bind_params)


def inject_view_propagation(sharding_info):

    view_op_list = ["reshape"]

    for op_name in sharding_info:
        if op_name in view_op_list:
            for args in sharding_info[op_name]:
                # example args for view:
                # "[((4, 256, 768), dtype('float32'))]{'new_sizes': (1, 4, 256, 768), 'dimensions': None}"
                re_matches = re.findall(r'\((\d+,|\d+(?:, \d+)*)\)', args)
                input_shape = [int(num) for num in re_matches[0].strip(',').split(',')]
                output_shape = [int(num) for num in re_matches[1].strip(',').split(',')]
                sharding_info[op_name][args] = view_propagation(input_shape, 
                                                                output_shape, 
                                                                world_size=device_mesh_world_size())

    return sharding_info


class EDJaxShardingAnn:

    def __init__(self, jaxpr, use_cache=True, seed=42) -> None:
        self.jaxpr = jaxpr
        self.use_cache = use_cache
        self.shape_info = {}
        self.sharding_info = {}

        self.key = jax.random.PRNGKey(seed)

    def run(self, consts, *args) -> Any:
        # Mapping from variable -> value
        env = {}

        def read(var):
            # Literals are values baked into the Jaxpr
            if type(var) is jax.core.Literal:
                return var.val
            if isinstance(env[var], ShapedArray):
                self.key, subkey = jax.random.split(self.key)
                if env[var].dtype.name in ["float64", "float32", "float16"]:
                    return jax.random.normal(subkey, shape=env[var].shape, dtype=env[var].dtype)
                elif env[var].dtype.name in ["int32", "unint32", "int64", "uint64", "uint8"]:
                    return jax.random.randint(subkey,
                                              shape=env[var].shape,
                                              dtype=env[var].dtype,
                                              minval=1,
                                              maxval=8)
                elif env[var].dtype.name in ["bool"]:
                    return jax.random.normal(subkey, shape=env[var].shape) > 1.
                else:
                    return jax.numpy.empty(shape=env[var].shape, dtype=env[var].dtype)
            return env[var]

        def write(var, val):
            if isinstance(val, int):
                env[var] = ShapedArray(shape=set(), dtype=jax.numpy.int32)
            elif isinstance(val, jax.Array) and not jax.core.is_opaque_dtype(val.dtype):
                env[var] = ShapedArray(shape=val.shape, dtype=val.dtype)
            else:
                env[var] = val

        def record_shape(var):
            if isinstance(env[var], jax.Array) or isinstance(env[var], ShapedArray):
                self.shape_info[var.__str__()] = {
                    'shape': env[var].shape,
                    'dtype': env[var].dtype.name
                }

        flat_args, _ = jax.tree_util.tree_flatten(args)

        # Bind args and consts to environment
        safe_map(write, self.jaxpr.invars, flat_args)
        safe_map(write, self.jaxpr.constvars, consts)

        safe_map(record_shape, self.jaxpr.invars)
        safe_map(record_shape, self.jaxpr.constvars)

        # add sharding_info for invars+constvars
        for invar in self.jaxpr.invars + self.jaxpr.constvars:
            invals = env[invar]
            onelevel_key = "placeholder"
            twolevel_key = str([(invals.shape, invals.dtype.name)])
            if onelevel_key not in self.sharding_info:
                self.sharding_info[onelevel_key] = {}
            self.sharding_info[onelevel_key][twolevel_key] = view_propagation(invals.shape, 
                                                                              invals.shape, 
                                                                              world_size=device_mesh_world_size())

        # Loop through equations and evaluate primitives using `bind`
        for eqn in self.jaxpr.eqns:
            subfuns, bind_params = eqn.primitive.get_bind_params(eqn.params)
            onelevel_key = eqn.primitive.__str__()
            name_stack = source_info_util.current_name_stack() + eqn.source_info.name_stack

            with source_info_util.user_context(eqn.source_info.traceback, name_stack=name_stack):

                # Read inputs to equation from environment
                invals = safe_map(read, eqn.invars)

                if eqn.primitive.__str__() == "custom_jvp_call":
                    pjit_eqn = subfuns[0].f.args[0].eqns[0]
                    onelevel_key += "[" + pjit_eqn.params['name'] + "]"
                elif eqn.primitive.__str__() in ["xla_call", "pjit"]:
                    onelevel_key += "[" + eqn.params['name'] + "]"

                # `bind` is how a primitive is called
                meta_op_ = MetaOp(func=eqn.primitive.bind,
                                  input_args=(subfuns, invals, bind_params),
                                  name=onelevel_key)
                outvals = meta_op_.exec()

            if onelevel_key not in self.sharding_info:
                self.sharding_info[onelevel_key] = {}

            twolevel_key = generate_str(invals, bind_params)

            if not self.use_cache or twolevel_key not in self.sharding_info[onelevel_key]:
                prompt_annotation = None
                if len(self.sharding_info[onelevel_key]) >= 1:
                    prompt_annotation = list(
                        self.sharding_info[onelevel_key].values())[0]["sharding_ann"]
                sharding_ann, combination_ann = meta_op_.sharding_discovery(
                    prompt_annotation=prompt_annotation)
                self.sharding_info[onelevel_key][twolevel_key] = {
                    "sharding_ann": sharding_ann,
                    "combination_ann": combination_ann
                }

            # outvals = eqn.primitive.bind(*invals, **eqn.params)
            # Primitives may return multiple outputs or not
            if not eqn.primitive.multiple_results:
                outvals = [outvals]
            # Write the results of the primitive into the environment
            safe_map(write, eqn.outvars, outvals)
            safe_map(record_shape, eqn.outvars)

        # Read the final result of the Jaxpr from the environment
        # return safe_map(read, self.jaxpr.outvars)
        return inject_view_propagation(self.sharding_info), self.shape_info

easydist/jax/sharding_interpreter.py
alibaba/easydist
# Copyright (c) 2023, Alibaba Group;
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================

import os
from contextlib import contextmanager


@contextmanager
def _sharding_ann_env():

    ori_tf32_override = os.environ.get("NVIDIA_TF32_OVERRIDE", None)

    os.environ["NVIDIA_TF32_OVERRIDE"] = "1"

    try:
        yield
    finally:
        if ori_tf32_override is None:
            os.environ.pop("NVIDIA_TF32_OVERRIDE")
        os.environ["NVIDIA_TF32_OVERRIDE"] = ori_tf32_override

easydist/jax/utils.py
alibaba/easydist
# Copyright (c) 2023, Alibaba Group;
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================

from .metaop import MetaOp
from .view_propagation import view_propagation
from .annotation import ShardDim, ShardAnnotation

__all__ = ["MetaOp", "view_propagation", "ShardDim", "ShardAnnotation"]

easydist/metashard/__init__.py
alibaba/easydist
# Copyright (c) 2023, Alibaba Group;
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================

import copy
from typing import List

from easydist import platform
from easydist.metashard.halo import HaloInfo


class ShardDim:
    """
    ``ShardDim`` is used to store sharding information for ``ShardAnnotation``.

    Within ``ShardAnnotation``, each dimension of the tensor list is assigned a ``ShardDim``. 
    A value of ``ShardDim(0)`` means that the dimension cannot be sharded. 
    All dimensions with a value of ``ShardDim(i)`` can be sharded together.

    ShardDim can also carry information about halo and chunk.

    - ``halo`` means that padding is needed for each shard.
    - ``chunk`` means that we need firstly chunk the data along the dimension,
      then shard for each chunk to get the sharded data. similar to block-cyclic.

      [1, 1, 2, 2] -(shard)-> [1, 1] | [2, 2] when chunk = 1
      [1, 1, 2, 2] -(chunk)-> [1, 1] and [2, 2] -(shard)-> [1, 2] | [1, 2]

    """

    def __init__(self, shard_dim_id: int, chunk: int = 1) -> None:
        self.shard_dim_id = shard_dim_id
        self.chunk = chunk

        self.halo: HaloInfo = None

    def __str__(self) -> str:
        if self.shard_dim_id == 0:
            return "NoShardDim"
        else:
            content = str(self.shard_dim_id)
            if self.chunk > 1:
                content += f", chunk={self.chunk})"
            if self.halo:
                content += f", halo={self.halo})"
            return f"ShardDim({content})"

    def set_halo(self, halo_: HaloInfo):
        self.halo = halo_

    def __repr__(self) -> str:
        return self.__str__()

    @staticmethod
    def get_noshard_dim():
        return ShardDim(0)

    @staticmethod
    def get_shard_dim(shard_dim_id):
        return ShardDim(shard_dim_id)


class ShardAnnotation:
    """
    ``ShardAnnotation`` used to describe the sharding strategy space for a ``MetaOp``.
    
    For example:
    - matmul: ShardAnnotation([[ShardDim(1), ShardDim(2)], [ShardDim(2), ShardDim(3)]])
    - relu: ShardAnnotation([[ShardDim(1), ShardDim(2)], [ShardDim(1), ShardDim(2)]])
    - layernorm: ShardAnnotation([[ShardDim(1), ShardDim(2), NoShardDim]])
    """

    def __init__(self, annotation: List[List[ShardDim]]) -> None:
        self.annotation = annotation

    @staticmethod
    def init_from_input_args(input_args):
        return ShardAnnotation([[ShardDim.get_noshard_dim()] * len(t.shape) for t in input_args
                                if isinstance(t, platform.Tensor)])

    def inject_haloinfo(self, haloinfo: HaloInfo, shard_dim_idx: int):
        """set haloinfo for each ShardDim(shard_dim_idx) in ShardAnnotation"""
        if haloinfo is None:
            return
        for anno in self.annotation:
            for shard_dim in anno:
                if shard_dim.shard_dim_id == shard_dim_idx:
                    shard_dim.set_halo(haloinfo)

    def get_max_shard_dim_id(self) -> int:
        max_shard_dim_id = 0
        for anno in self.annotation:
            for dim_anno in anno:
                max_shard_dim_id = max(max_shard_dim_id, dim_anno.shard_dim_id)
        return max_shard_dim_id

    def clear_shard_dim(self, max_shard_dim_id):
        new_anno = copy.deepcopy(self)
        for anno in new_anno.annotation:
            for dim_idx in range(len(anno)):
                if anno[dim_idx].shard_dim_id > max_shard_dim_id:
                    anno[dim_idx] = ShardDim.get_noshard_dim()
        return new_anno

    def __str__(self) -> str:
        return f"ShardAnnotation({self.annotation.__str__()})"

    def __repr__(self) -> str:
        return f"ShardAnnotation({self.annotation.__repr__()})"

    def __getitem__(self, idx):
        return self.annotation[idx]

    def __setitem__(self, idx, shard_dim):
        self.annotation[idx] = shard_dim

    def __len__(self) -> int:
        return self.annotation.__len__()

    def __add__(self, other):
        return ShardAnnotation(self.annotation + other.annotation)

easydist/metashard/annotation.py
alibaba/easydist
# Copyright (c) 2023, Alibaba Group;
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================

import functools
from enum import Enum
import logging

from easydist import platform
import easydist.config as mdconfig

logger = logging.getLogger(__name__)

TRY_COMBINATION_FUNC = dict()


def try_combination_func(func):
    """Register a function as a try_combination_func"""
    TRY_COMBINATION_FUNC[func.__name__] = func
    return func


class ReduceOp(Enum):
    SUM = 1
    MAX = 2
    MIN = 3


class HaloHint:

    def __init__(self, halo, dim, idx_=None):
        self.halo = halo
        self.dim = dim
        self.out_idx = idx_


def aligned_prefix(t1: platform.Tensor, t2: platform.Tensor, dim_idx: int):
    """calculate the max prefix for t1 and t2 on dimension(dim_idx)

    For example: aligned_prefix([1, 2, 3], [1, 2, 5]) -> 2
    """
    dim_size = min(t1.shape[dim_idx], t2.shape[dim_idx])
    for i in range(1, dim_size + 1):
        if not platform.allclose(platform.narrow(t1, dim_idx, 0, i),
                                 platform.narrow(t2, dim_idx, 0, i)):
            return i - 1
    return i


def shape_aligned_otherdim(shape_1, shape_2, dim_idx):
    """return True when shape_1 and shape_2 only different on dimension(dim_idx)"""
    if len(shape_1) != len(shape_2):
        return False

    diff_dim = []
    for idx in range(len(shape_1)):
        if shape_1[idx] != shape_2[idx]:
            diff_dim.append(idx)

    if diff_dim == [dim_idx]:
        return True
    return False


class CombinationFunc:

    @staticmethod
    def identity(sharded_tensor):
        identity_tensor = sharded_tensor[0]
        for tensor_ in sharded_tensor:
            if not platform.equal(identity_tensor, tensor_):
                logger.debug("not all tensor same as identity")
                return None
        return identity_tensor

    @staticmethod
    def reduce(sharded_tensor, ops=ReduceOp.SUM):
        init = platform.zeros_like(sharded_tensor[0])
        assert ops in [ReduceOp.SUM, ReduceOp.MAX, ReduceOp.MIN]
        if ops == ReduceOp.SUM:
            reduce_func_ = platform.add
        elif ops == ReduceOp.MAX:
            reduce_func_ = platform.max
        elif ops == ReduceOp.MIN:
            reduce_func_ = platform.min
        return functools.reduce(reduce_func_, sharded_tensor, init)

    @staticmethod
    def gather(sharded_tensor, dim, halowidth=0, chunk=1):
        if halowidth == 0:
            if chunk == 1:
                return platform.concatenate(sharded_tensor, dim=dim)
            else:
                chunk_sharded_tensor = [platform.chunk(t, chunk, dim) for t in sharded_tensor]
                reorder_list = []
                for chunk_idx in range(chunk):
                    reorder_list += [chunk_list[chunk_idx] for chunk_list in chunk_sharded_tensor]
                return platform.concatenate(reorder_list, dim=dim)

        shard_world_size = len(sharded_tensor)
        gathered_tensor = sharded_tensor[0]
        for idx_ in range(1, shard_world_size):
            width_tensor_1 = int(gathered_tensor.shape[dim])
            width_tensor_2 = int(sharded_tensor[idx_].shape[dim])
            if halowidth > 0:
                gathered_tensor = platform.concatenate([
                    platform.narrow(gathered_tensor, dim, 0, width_tensor_1 - halowidth),
                    platform.add(
                        platform.narrow(gathered_tensor, dim, width_tensor_1 - halowidth,
                                        halowidth),
                        platform.narrow(sharded_tensor[idx_], dim, 0, halowidth)),
                    platform.narrow(sharded_tensor[idx_], dim, halowidth,
                                    width_tensor_2 - halowidth)
                ],
                                                       dim=dim)
            else:
                gathered_tensor = platform.concatenate([
                    platform.narrow(gathered_tensor, dim, 0, width_tensor_1 + halowidth),
                    platform.narrow(sharded_tensor[idx_], dim, -1 * halowidth,
                                    width_tensor_2 + halowidth)
                ],
                                                       dim=dim)

        return gathered_tensor


@try_combination_func
def try_combination_identity(sharded_tensor_, global_tensor_):
    """try combination through identity, return combination func if success, otherwise None"""

    for t_ in sharded_tensor_:
        if t_.shape != global_tensor_.shape:
            return None

    result = CombinationFunc.identity(sharded_tensor_)
    if result is not None and platform.allclose(result, global_tensor_):
        return functools.partial(CombinationFunc.identity)

    return None


@try_combination_func
def try_combination_reduce(sharded_tensor_, global_tensor_):
    """try combination through reduce, return combination func if success, otherwise None"""

    for t_ in sharded_tensor_:
        if t_.shape != global_tensor_.shape:
            return None

    for reduce_op in [ReduceOp.SUM, ReduceOp.MAX, ReduceOp.MIN]:
        reduce_func = functools.partial(CombinationFunc.reduce, ops=reduce_op)
        if platform.allclose(reduce_func(sharded_tensor_), global_tensor_):
            return reduce_func

    return None


@try_combination_func
def try_combination_gather(sharded_output_, global_output_):
    """try combination through gather, return combination func if success, otherwise None"""

    # 1) all sharded tensor shape should be equal to global tensor except the gather_dim
    sharded_output_shape = sharded_output_[0].shape
    global_output_shape = global_output_.shape

    # return if global_output is only one element.
    if len(global_output_shape) == 0:
        return None

    shard_size = len(sharded_output_)

    for dim_idx in range(len(sharded_output_shape)):
        if sharded_output_shape[dim_idx] != global_output_shape[dim_idx]:
            break

    for tensor_ in sharded_output_:
        if not shape_aligned_otherdim(tensor_.shape, global_output_shape, dim_idx):
            return None

    combination_size = 0
    for out_ in sharded_output_:
        combination_size += out_.shape[dim_idx]
    gap_size = combination_size - global_output_shape[dim_idx]
    gather_func = None
    if gap_size == 0:
        gather_func = functools.partial(CombinationFunc.gather, dim=dim_idx)
        gathered_tensor = gather_func(sharded_output_)
        if gathered_tensor.shape == global_output_.shape:
            if platform.allclose(gathered_tensor, global_output_):
                return gather_func
            if mdconfig.extend_space:
                reference_shard_output = platform.chunk(global_output_,
                                                        chunks=shard_size,
                                                        dim=dim_idx)[0]
                aligned_prefix_len = aligned_prefix(sharded_output_[0], reference_shard_output,
                                                    dim_idx)

                # explore gather with chunk > 1
                if aligned_prefix_len != 0 and sharded_output_shape[
                        dim_idx] % aligned_prefix_len == 0:
                    guess_chunk = sharded_output_shape[dim_idx] // aligned_prefix_len
                    gather_func = functools.partial(CombinationFunc.gather,
                                                    dim=dim_idx,
                                                    chunk=guess_chunk)
                    gathered_tensor_chunk = gather_func(sharded_output_)
                    if platform.allclose(gathered_tensor_chunk, global_output_):
                        return gather_func

                # return HaloHint to explore the chance of halo sharding
                if aligned_prefix_len > (sharded_output_shape[dim_idx] // 2):
                    return HaloHint(sharded_output_shape[dim_idx] - aligned_prefix_len, dim_idx)

    if mdconfig.extend_space:
        # halo gather with halowidth > 0
        if gap_size > 0 and gap_size % (shard_size - 1) == 0:
            halowidth = int(gap_size // (shard_size - 1))
            if halowidth >= combination_size // shard_size:
                return None
            gather_func = functools.partial(CombinationFunc.gather,
                                            dim=dim_idx,
                                            halowidth=halowidth)
            gathered_tensor = gather_func(sharded_output_)
            if gathered_tensor.shape == global_output_.shape:
                if platform.allclose(gathered_tensor, global_output_):
                    return gather_func

        # halo gather with halowidth < 0
        if gap_size > 0 and gap_size % shard_size == 0:
            halowidth = -1 * int(gap_size // shard_size)
            if -1 * halowidth >= combination_size // (2 * shard_size):
                return None
            gather_func = functools.partial(CombinationFunc.gather,
                                            dim=dim_idx,
                                            halowidth=halowidth)
            gathered_tensor = gather_func(sharded_output_)
            if gathered_tensor.shape == global_output_.shape:
                if platform.allclose(gathered_tensor, global_output_):
                    return gather_func

        # raise HaloHint when convolution without padding
        if gap_size < 0 and gap_size % (shard_size - 1) == 0:
            halowidth = int(gap_size // (shard_size - 1)) // 2
            if -1 * halowidth >= combination_size // shard_size:
                return None
            return HaloHint(halowidth, dim_idx)


def try_combination_single(sharded_output_, global_output_):

    # check all sharded tensor have equal dimension of global_output
    for sharded_tensor in sharded_output_:
        if len(sharded_tensor.shape) != len(global_output_.shape):
            return None

    for func_name in TRY_COMBINATION_FUNC:
        combination_func = TRY_COMBINATION_FUNC[func_name](sharded_output_, global_output_)
        if combination_func:
            return combination_func

    return None


def try_combination(sharded_output_, global_output_):

    if isinstance(global_output_, platform.Tensor):
        return try_combination_single(sharded_output_, global_output_)

    if isinstance(global_output_, tuple) or isinstance(global_output_, list):
        output_num = [len(i) for i in sharded_output_]
        if len(global_output_) == min(output_num) == max(output_num):
            return_combination_ann = []
            for idx_ in range(len(global_output_)):
                if isinstance(global_output_[idx_], platform.Tensor):
                    single_ann = try_combination_single([i[idx_] for i in sharded_output_],
                                                        global_output_[idx_])
                    if single_ann is None:
                        return None
                    if isinstance(single_ann, HaloHint):
                        single_ann.out_idx = idx_
                        return single_ann
                    return_combination_ann.append(single_ann)
                else:
                    for out_ in sharded_output_:
                        if global_output_[idx_] != out_[idx_]:
                            return None

            if len(return_combination_ann) > 0:
                return return_combination_ann

    return None

easydist/metashard/combination.py
alibaba/easydist
# Copyright (c) 2023, Alibaba Group;
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================

from typing import List

from easydist import platform


class HaloInfo:

    def __init__(self, halowidth: int, dim: int) -> None:
        self.halowidth = halowidth
        self.dim = dim

    def __str__(self) -> str:
        return self.halowidth.__str__()

    def __repr__(self) -> str:
        return self.__str__()


def halo_padding(tensor_list_: List[platform.Tensor], haloinfo: HaloInfo) -> List[platform.Tensor]:
    """add halo padding to tensor_list_"""

    if haloinfo is None or len(tensor_list_) < 2:
        return tensor_list_

    halo = haloinfo.halowidth
    dim = haloinfo.dim

    padded_tensor_list = []
    for idx in range(len(tensor_list_)):
        to_concatenate = [tensor_list_[idx]]
        if idx >= 1:
            dim_size = tensor_list_[idx - 1].shape[dim]
            if dim_size < halo:
                raise RuntimeError("Cannot halo padding for this sharded_tensor")
            to_concatenate.insert(
                0, platform.narrow(tensor_list_[idx - 1], dim, dim_size - halo, halo))
        if idx <= len(tensor_list_) - 2:
            to_concatenate.append(platform.narrow(tensor_list_[idx + 1], dim, 0, halo))
        padded_tensor_list.append(platform.concatenate(to_concatenate, dim=dim))

    return padded_tensor_list

easydist/metashard/halo.py
alibaba/easydist
# Copyright (c) 2023, Alibaba Group;
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================

from __future__ import annotations
import copy
import os
import logging
from typing import List
from functools import reduce

from easydist.metashard.combination import ReduceOp
from easydist.platform import get_backend
import easydist.config as mdconfig

logger = logging.getLogger(__name__)

DEVICE_MESH_1D = -1


class SPMD:

    REPLICATE = "REPLICATE"
    SHARD = "SHARD"
    PARTIAL = "PARTIAL"

    def __init__(self, state, args=None) -> None:
        self.state = state
        self.args = args

    def is_shard(self) -> bool:
        return self.state == SPMD.SHARD

    def is_replicate(self) -> bool:
        return self.state == SPMD.REPLICATE

    def is_partial(self) -> bool:
        return self.state == SPMD.PARTIAL

    def __eq__(self, other) -> bool:
        if self.state == other.state and self.args == other.args:
            return True
        return False

    def __str__(self) -> str:
        str_ = f"{self.state}"
        if self.args:
            str_ += f"({self.args})"
        return str_

    def __repr__(self) -> str:
        return self.__str__()


class VarSPMDStrategy:

    def __init__(self, *var_spmd_strategy) -> None:
        self.var_spmd_strategy = list(var_spmd_strategy)

    def __getitem__(self, idx) -> SPMD:
        return self.var_spmd_strategy[idx]

    def __add__(self, other):
        return VarSPMDStrategy(*self.var_spmd_strategy, *other.var_spmd_strategy)

    def __eq__(self, other) -> bool:
        if len(self.var_spmd_strategy) == len(other.var_spmd_strategy):
            for i, j in zip(self.var_spmd_strategy, other.var_spmd_strategy):
                if i != j:
                    return False
            return True
        return False

    def __str__(self) -> str:
        return f"VarSPMDStrategy({self.var_spmd_strategy.__str__()})"

    def __len__(self) -> str:
        return self.var_spmd_strategy.__len__()

    def __repr__(self) -> str:
        return self.__str__()


class VarSPMDStrategyGroup:

    def __init__(self, *var_spmd_strategy_group) -> None:
        self.var_spmd_strategy_group = list(var_spmd_strategy_group)

    def append(self, var_spmd_strategy: VarSPMDStrategy) -> None:
        self.var_spmd_strategy_group.append(var_spmd_strategy)

    def get_var_strtg(self, idx: int) -> VarSPMDStrategy:
        assert idx < len(self.var_spmd_strategy_group)
        return self.var_spmd_strategy_group[idx]

    def __eq__(self, other) -> bool:
        if len(self.var_spmd_strategy_group) == len(other.var_spmd_strategy_group):
            for i, j in zip(self.var_spmd_strategy_group, other.var_spmd_strategy_group):
                if i != j:
                    return False
            return True
        return False

    def __getitem__(self, idx) -> VarSPMDStrategy:
        return self.var_spmd_strategy_group[idx]

    def __len__(self) -> str:
        return self.var_spmd_strategy_group.__len__()

    def __str__(self) -> str:
        return f"VarSPMDStrategyGroup({self.var_spmd_strategy_group.__str__()})"

    def __repr__(self) -> str:
        return self.__str__()


class NodeSPMDStrategy:

    def __init__(self, in_strtg_group: VarSPMDStrategyGroup,
                 out_strtg_group: VarSPMDStrategyGroup):
        self.in_strtg_group = in_strtg_group
        self.out_strtg_group = out_strtg_group

    def get_invar_strtg(self, invar_idx: int) -> VarSPMDStrategy:
        return self.in_strtg_group.get_var_strtg(invar_idx)

    def get_outvar_strtg(self, outvar_idx: int) -> VarSPMDStrategy:
        return self.out_strtg_group.get_var_strtg(outvar_idx)

    def __str__(self) -> str:
        return (f"NodeSPMDStrategy(in_strtg_group: {self.in_strtg_group}, "
                f"out_strtg_group: {self.out_strtg_group})")

    def __repr__(self) -> str:
        return self.__str__()


class NodeSPMDStrategyPool:

    def __init__(self):
        self.strategies = []

    def add_strategy(self, strtg: NodeSPMDStrategy):
        self.strategies.append(strtg)

    def find_matched_out(self, out_var_idx: int, expected_strtg: VarSPMDStrategy):
        for strtg_idx, nd_strtg in enumerate(self.strategies):
            if nd_strtg.out_strtg_group.get_var_strtg(out_var_idx) == expected_strtg:
                return strtg_idx

        return -1

    def get_strtg(self, idx: int) -> NodeSPMDStrategy:
        assert idx < len(self.strategies)
        return self.strategies[idx]

    def strtg_num(self) -> int:
        return len(self.strategies)

    def __str__(self) -> str:
        return f"{self.strategies})"

    def __repr__(self) -> str:
        return self.__str__()


def get_sharding_strategy(sharding_anns, shard_dim_id):
    spmd_strategy = VarSPMDStrategyGroup()

    for tensor_ann in sharding_anns.annotation:
        shard_dim_ = None
        for dim_idx, dim_ann in enumerate(tensor_ann):
            if dim_ann.shard_dim_id == shard_dim_id:
                shard_dim_ = dim_idx
        if shard_dim_ is not None:
            spmd_strategy.append(VarSPMDStrategy(SPMD(SPMD.SHARD, {"dim": shard_dim_})))
        else:
            spmd_strategy.append(VarSPMDStrategy(SPMD(SPMD.REPLICATE)))

    return spmd_strategy


def combination_to_sharding_strategy(comm_anns, all_replicate=False):
    # example of comm_anns:
    # functools.partial(<function CombinationFunc.gather at 0x7fab788efd30>, dim=0)

    spmd_strategy = VarSPMDStrategyGroup()

    if not (isinstance(comm_anns, list) or isinstance(comm_anns, tuple)):
        comm_anns = [comm_anns]
    for comm_ann in comm_anns:
        func_name = comm_ann.func.__name__
        if all_replicate or func_name == "identity":
            spmd_strategy.append(VarSPMDStrategy(SPMD(SPMD.REPLICATE)))
        elif func_name == "gather":
            spmd_strategy.append(VarSPMDStrategy(SPMD(SPMD.SHARD, comm_ann.keywords)))
        elif func_name == "reduce":
            spmd_strategy.append(VarSPMDStrategy(SPMD(SPMD.PARTIAL, comm_ann.keywords)))

    return spmd_strategy


_dtype2byte = {
    "float64": 8,
    "float32": 4,
    "float16": 2,
    "bool": 0.125,
    "int32": 4,
    "int64": 8,
    "uint32": 4,
    "uint8": 1,
    "complex64": 16,
}


class MetaVar:

    id_counter: int = 0

    @staticmethod
    def generate_unique_id():
        unique_id = MetaVar.id_counter
        MetaVar.id_counter += 1
        return unique_id

    @staticmethod
    def clear_id_counter():
        MetaVar.id_counter = 0

    def __init__(self, name, shape, dtype) -> None:
        self.unique_id = self.generate_unique_id()
        self.name = name
        self.shape = shape
        self.dtype = dtype
        self.up_node = None
        self.idx_for_up = -1
        self.down_nodes = []
        self.indice_for_down = []

    def get_var_size(self):
        if len(self.shape) == 0:
            return _dtype2byte[self.dtype]
        num_ele = reduce((lambda x, y: x * y), self.shape)
        return _dtype2byte[self.dtype] * num_ele

    def __str__(self, details=False) -> str:
        return self.name.__str__()

    def debug_str(self) -> str:
        return f"{self.name}({self.shape}, {self.dtype})"

    def __repr__(self) -> str:
        return self.__str__()


# (TODO) need to automatic find the heavy ops during MetaSPMD annotation.
_heavy_ops = [
    "convolution",
    "scaled_dot_product",
    "aten.baddbmm",
    "aten.mm",
    "aten.bmm",
    "dot",
]


class MetaNode:

    id_counter: int = 0

    @staticmethod
    def generate_unique_id():
        unique_id = MetaVar.id_counter
        MetaVar.id_counter += 1
        return unique_id

    @staticmethod
    def clear_id_counter():
        MetaVar.id_counter = 0

    def __init__(self, name, op_name, invars: List[MetaVar], outvars: List[MetaVar],
                 sharding_info, is_placeholder=False) -> None:
        self.unique_id = self.generate_unique_id()
        self.cluster_id = -1
        self.name = name
        self.op_name = op_name
        self.invars = invars
        for idx, in_var in enumerate(invars):
            if in_var:
                if isinstance(in_var, MetaVar):  # Lansong(TODO): should remove this line later
                    in_var.down_nodes.append(self)
                    in_var.indice_for_down.append(idx)
        self.compact_out_idx_tbl = []

        self.outvars = outvars
        for idx, out_var in enumerate(outvars):
            if out_var:
                out_var.up_node = self
                out_var.idx_for_up = idx
        self.compact_in_idx_tbl = []

        self.strtg_pool = None
        self.unique_key_ = None

        self.sharding_info = sharding_info
        self.is_placeholder = is_placeholder

    def set_in_var(self, in_var: MetaVar, idx: int):
        assert idx < len(self.compact_in_idx_tbl)
        compact_idx = self.compact_in_idx_tbl[idx]

        assert compact_idx < len(self.invars)
        assert not self.invars[compact_idx]
        self.invars[compact_idx] = in_var
        in_var.down_nodes.append(self)
        in_var.indice_for_down.append(compact_idx)

    def set_out_var(self, out_var: MetaVar, idx: int):
        assert idx < len(self.compact_out_idx_tbl)
        compact_idx = self.compact_out_idx_tbl[idx]

        assert compact_idx < len(self.outvars)
        assert not self.outvars[compact_idx]
        self.outvars[compact_idx] = out_var
        out_var.up_node = self
        out_var.idx_for_up = compact_idx

    def unique_key(self):
        if self.unique_key_ is None:
            self.unique_key_ = str(self.name)
        return self.unique_key_

    def _replicate_strategy(self):
        invars_strategy = VarSPMDStrategyGroup()
        for _ in self.invars:
            invars_strategy.append(VarSPMDStrategy(SPMD(SPMD.REPLICATE)))
        outvars_strategy = VarSPMDStrategyGroup()
        for _ in self.outvars:
            outvars_strategy.append(VarSPMDStrategy(SPMD(SPMD.REPLICATE)))
        return NodeSPMDStrategy(invars_strategy, outvars_strategy)

    def get_strtg(self, idx: int) -> NodeSPMDStrategy:
        assert self.strtg_pool
        return self.strtg_pool.get_strtg(idx)

    def get_strtg_pool(self) -> NodeSPMDStrategyPool:
        if self.strtg_pool is not None:
            return self.strtg_pool

        strategy_list_1d = []

        sharding_anns = self.sharding_info['sharding_ann']
        comm_anns = self.sharding_info['combination_ann']

        # some ops like torch.ops.aten.scalar_tensor.default have no invars
        if len(comm_anns) == 0 and len(self.invars) == 0:
            self.strtg_pool = NodeSPMDStrategyPool()
            return self.strtg_pool

        for comm_ann in comm_anns:
            in_strtg_group = get_sharding_strategy(sharding_anns, comm_ann)
            out_strtg_group = combination_to_sharding_strategy(comm_anns[comm_ann])
            strategy_list_1d.append(NodeSPMDStrategy(in_strtg_group, out_strtg_group))

        if all(op not in self.op_name for op in _heavy_ops) or len(comm_anns) == 0:
            strategy_list_1d.append(self._replicate_strategy())

        # (FIXME) modify stratgy here, need remove
        #     if "batch_norm" in self.name:
        #         dp_strategy = {
        #             "invars_sharding":
        #             get_sharding_strategy(sharding_anns, None),
        #             "outvars_sharding":
        #             combination_to_sharding_strategy(comm_anns[comm_ann], all_replicate=True)
        #         }
        #         dp_strategy["invars_sharding"][0] = SPMD(SPMD.SHARD, {'dim': 0})
        #         dp_strategy["outvars_sharding"][0] = SPMD(SPMD.SHARD, {'dim': 0})
        #         if "backward" in self.name:
        #             dp_strategy["invars_sharding"][1] = SPMD(SPMD.SHARD, {'dim': 0})
        #             dp_strategy["outvars_sharding"][1] = SPMD(SPMD.PARTIAL, {'ops': ReduceOp.SUM})
        #             dp_strategy["outvars_sharding"][2] = SPMD(SPMD.PARTIAL, {'ops': ReduceOp.SUM})
        #         strategy_list = [dp_strategy, replicate_strategy]

        # if "convolution.default" in self.name:
        #     del strategy_list[1]
        # if "convolution_backward.default" in self.name:
        #     del_idx = None
        #     for idx, s in enumerate(strategy_list):
        #         if s["invars_sharding"][-1] == SPMD(SPMD.SHARD, {'dim': 1}):
        #             del_idx = idx
        #             break
        #     if del_idx is not None:
        #         del strategy_list[del_idx]

        self.strtg_pool = NodeSPMDStrategyPool()
        if DEVICE_MESH_1D == -1:
            for idx1, s1 in enumerate(strategy_list_1d):
                for idx2, s2 in enumerate(strategy_list_1d):
                    # [Shard(i), Shard(i)] is not support for pytorch dtensor runtime
                    if get_backend() == "torch":
                        if any(i.state == SPMD.SHARD for i in s1.in_strtg_group) and idx1 == idx2:
                            continue
                    invars_strategy = VarSPMDStrategyGroup()
                    for i, j in zip(s1.in_strtg_group, s2.in_strtg_group):
                        invars_strategy.append(i + j)
                    outvars_strategy = VarSPMDStrategyGroup()
                    for i, j in zip(s1.out_strtg_group, s2.out_strtg_group):
                        outvars_strategy.append(i + j)
                    self.strtg_pool.add_strategy(
                        NodeSPMDStrategy(invars_strategy, outvars_strategy))
        else:
            replicate_ = VarSPMDStrategy(SPMD(SPMD.REPLICATE))
            for s in strategy_list_1d:
                if DEVICE_MESH_1D == 0:
                    invars_strategy = VarSPMDStrategyGroup()
                    for i in s.in_strtg_group:
                        invars_strategy.append(replicate_ + i)
                    outvars_strategy = VarSPMDStrategyGroup()
                    for i in s.out_strtg_group:
                        outvars_strategy.append(replicate_ + i)
                    self.strtg_pool.add_strategy(
                        NodeSPMDStrategy(invars_strategy, outvars_strategy))
                elif DEVICE_MESH_1D == 1:
                    invars_strategy = VarSPMDStrategyGroup()
                    for i in s.in_strtg_group:
                        invars_strategy.append(i + replicate_)
                    outvars_strategy = VarSPMDStrategyGroup()
                    for i in s.out_strtg_group:
                        outvars_strategy.append(i + replicate_)
                    self.strtg_pool.add_strategy(
                        NodeSPMDStrategy(invars_strategy, outvars_strategy))
                else:
                    exit(-1)

        if mdconfig.log_level <= logging.DEBUG:
            print("node: %s" % self.name.__str__())
            print(self.strtg_pool)
        return self.strtg_pool

    def __str__(self) -> str:
        return self.name.__str__()

    def debug_str(self) -> str:
        return_str = f"{self.name} ({self.op_name})\n"
        return_str += "invars: " + ",".join([var.debug_str() for var in self.invars]) + "\n"
        return_str += "outvars: " + ",".join(
            [var.debug_str() if var is not None else "None" for var in self.outvars]) + "\n"
        return_str += f"sharding_info: {self.sharding_info}"
        return return_str

    def __repr__(self) -> str:
        return self.__str__()


class ClusterStrategy:

    def __init__(self):
        self.in_strtg_group = None
        self.out_strtg_group = None
        self.inner_strtg_group = None
        self.node_strategies = {}

    def set_node_strategy(self, nd_id: int, nd_strtg_id: int, nd_strtg: NodeSPMDStrategy):
        assert nd_id not in self.node_strategies

        self.node_strategies[nd_id] = [nd_strtg_id, nd_strtg]

    def __str__(self) -> str:
        return (f"ClusterStrategies(in_strtg_group: {self.in_strtg_group}, "
                f"out_strtg_group: {self.out_strtg_group}, "
                f"inner_strtg_group: {self.inner_strtg_group})")

    def __repr__(self) -> str:
        return self.__str__()


class NodeIOStrategies:

    def __init__(self, node: MetaNode):
        self.node = node
        self.in_strategies = [[] for _ in range(len(node.invars))]
        if self.node.is_placeholder:
            self.in_strategies.append([])
        self.out_strategies = [[] for _ in range(len(node.outvars))]

    def add_in_strategy(self, in_idx: int, var_strtg: VarSPMDStrategy):
        if not self.node.is_placeholder:
            assert in_idx < len(self.node.invars)
        self.in_strategies[in_idx].append(var_strtg)

    def add_out_strategy(self, out_idx: int, var_strtg: VarSPMDStrategy):
        assert out_idx < len(self.node.outvars)
        self.out_strategies[out_idx].append(var_strtg)

    def get_invar_strtg(self, invar_idx: int, strtg_idx: int) -> VarSPMDStrategy:
        return self.in_strategies[invar_idx][strtg_idx]

    def get_outvar_strtg(self, outvar_idx: int, strtg_idx: int) -> VarSPMDStrategy:
        return self.out_strategies[outvar_idx][strtg_idx]

    def get_invar_strtg_list(self, invar_idx: int):  # return list of VarSPMDStrategy
        return self.in_strategies[invar_idx]

    def get_outvar_strtg_list(self, outvar_idx: int):  # return list of VarSPMDStrategy
        return self.out_strategies[outvar_idx]

    def __str__(self) -> str:
        res = str(self.node)
        res += f"\nin strategies: {self.in_strategies}"
        res += f"\noutput strategies: {self.out_strategies}"
        return res

    def __repr__(self) -> str:
        return self.__str__()


class ClusterStrategyPool:

    def __init__(self, cluster: MetaNodeCluster):
        self.cluster = cluster
        self.node_io_strategies = {}
        self.node_strategies = {}
        self.strtg_num = 0

        for nd in cluster.nodes.values():
            self.node_io_strategies[nd.unique_id] = NodeIOStrategies(nd)
            self.node_strategies[nd.unique_id] = []

    def get_invar_strtg(self, nd_id: int, invar_idx: int, strtg_idx: int) -> VarSPMDStrategy:
        assert strtg_idx < self.strtg_num
        nd_io_strtg = self.node_io_strategies[nd_id]
        return nd_io_strtg.get_invar_strtg(invar_idx, strtg_idx)

    def get_outvar_strtg(self, nd_id: int, outvar_idx: int, strtg_idx: int) -> VarSPMDStrategy:
        assert strtg_idx < self.strtg_num
        nd_io_strtg = self.node_io_strategies[nd_id]
        return nd_io_strtg.get_outvar_strtg(outvar_idx, strtg_idx)

    def get_invar_strtg_list(self, nd_id: int, invar_idx: int) -> list[VarSPMDStrategy]:
        nd_io_strtg = self.node_io_strategies[nd_id]
        return nd_io_strtg.get_invar_strtg_list(invar_idx)

    def get_outvar_strtg_list(self, nd_id: int, outvar_idx: int) -> list[VarSPMDStrategy]:
        nd_io_strtg = self.node_io_strategies[nd_id]
        return nd_io_strtg.get_outvar_strtg_list(outvar_idx)

    def get_node_strtg(self, nd_id: int, strtg_idx: int) -> NodeSPMDStrategy:
        assert nd_id in self.node_strategies
        assert strtg_idx < self.strtg_num
        return self.node_strategies[nd_id][strtg_idx]

    def add_strategy(self, strtg: ClusterStrategy):
        self.strtg_num += 1
        for nd_id, nd_strtg_info in strtg.node_strategies.items():
            nd_strtg_id = nd_strtg_info[0]
            nd_strtg = nd_strtg_info[1]
            nd = self.cluster.nodes[nd_id]

            # store node strategies with original format
            self.node_strategies[nd.unique_id].append(nd_strtg)

            # store io strategies
            nd_io_strtg = self.node_io_strategies[nd.unique_id]
            if nd.is_placeholder:
                outvar_strtg = nd_strtg.get_outvar_strtg(0)
                nd_io_strtg.add_in_strategy(0, outvar_strtg)
            else:
                for invar_idx in range(len(nd.invars)):
                    invar_strtg = nd_strtg.get_invar_strtg(invar_idx)
                    nd_io_strtg.add_in_strategy(invar_idx, invar_strtg)

                
            for outvar_idx in range(len(nd.outvars)):
                outvar_strtg = nd_strtg.get_outvar_strtg(outvar_idx)
                nd_io_strtg.add_out_strategy(outvar_idx, outvar_strtg)

    def get_strtg_num(self) -> int:
        return self.strtg_num

    def __str__(self) -> str:
        res = f"node strategies: {self.node_strategies}"
        res += f"\nnode io strategies: {self.node_io_strategies}"
        return res

    def __repr__(self) -> str:
        return self.__str__()


class ClusterArgs:

    def __init__(self) -> None:
        self.descs = []

    def add_arg(self, input_node: MetaNode, idx: int) -> None:
        self.descs.append([input_node, idx])

    def __str__(self) -> str:
        return_str = f"{self.descs.__str__()}"
        return return_str


class MetaNodeCluster:

    def __init__(self, unique_id: int) -> None:
        self.unique_id = unique_id
        self.args = ClusterArgs()
        self.output_node = None
        self.nodes = {}
        self.strategy_pool = None

    def add_node(self, meta_node: MetaNode) -> None:
        assert meta_node.unique_id not in self.nodes
        assert meta_node.unique_id >= 0
        self.nodes[meta_node.unique_id] = meta_node
        meta_node.cluster_id = self.unique_id

    def back_build_strategy(self, nd: MetaNode, nd_strtg_idx: int,
                            cluster_strtg: ClusterStrategy) -> bool:
        succ = True
        for invar_idx, invar in enumerate(nd.invars):
            if not invar:
                continue

            up_node = invar.up_node
            if not up_node:
                continue

            if up_node.unique_id not in self.nodes:
                # up_node is not in the cluster
                continue

            # we should guarantee no multi-output node inside cluster
            idx_for_up = invar.idx_for_up
            assert len(up_node.outvars) == 1
            assert idx_for_up == 0

            nd_strtg = nd.strtg_pool.get_strtg(nd_strtg_idx)
            expected_var_strtg = nd_strtg.get_invar_strtg(invar_idx)

            up_nd_strtg_pool = up_node.get_strtg_pool()
            up_nd_strtg_idx = up_nd_strtg_pool.find_matched_out(idx_for_up, expected_var_strtg)
            if up_nd_strtg_idx >= 0:
                up_nd_strtg = up_nd_strtg_pool.get_strtg(up_nd_strtg_idx)
                cluster_strtg.set_node_strategy(up_node.unique_id, up_nd_strtg_idx, up_nd_strtg)
                if not self.back_build_strategy(up_node, up_nd_strtg_idx, cluster_strtg):
                    succ = False
            else:
                succ = False

        return succ

    def finalize(self) -> None:
        for node in self.nodes.values():
            if node.invars:
                for idx, invar in enumerate(node.invars):
                    if not invar.up_node:
                        self.args.add_arg(node, idx)
                    elif invar.up_node.unique_id not in self.nodes:
                        self.args.add_arg(node, idx)
            else:
                # placeholder or get_attr node
                self.args.add_arg(node, 0)

        all_node_outs = []
        for node in self.nodes.values():
            all_node_outs.extend(node.outvars)

        # collect all variables which are read by outside of cluster
        cluster_out_vars = []
        for var in all_node_outs:
            if var:
                if not var.down_nodes:
                    # output MetaVar of graph has an empty down_nodes list
                    cluster_out_vars.append(var)
                else:
                    for node in var.down_nodes:
                        if node.unique_id not in self.nodes:
                            cluster_out_vars.append(var)

        # get or build strategy info for out node
        out_node = None
        for outvar in cluster_out_vars:
            if not out_node:
                out_node = outvar.up_node
            else:
                # be sure only one output node in a cluster
                assert out_node == outvar.up_node

        self.output_node = out_node

        # build strategy candidates for cluster
        self.strategy_pool = ClusterStrategyPool(self)
        out_strtg_pool = out_node.get_strtg_pool()

        for out_strtg_idx in range(out_strtg_pool.strtg_num()):
            cluster_strtg = ClusterStrategy()
            out_strtg = out_node.get_strtg(out_strtg_idx)
            cluster_strtg.set_node_strategy(out_node.unique_id, out_strtg_idx, out_strtg)

            if self.back_build_strategy(out_node, out_strtg_idx, cluster_strtg):
                assert len(cluster_strtg.node_strategies) == len(self.nodes)
                self.strategy_pool.add_strategy(cluster_strtg)
            else:
                # failed to search sync free cluster strategy
                # Lansong(TODO) we should decompose the cluster into multiple clusters
                logger.debug(
                    "failed to search sync free cluster strategy for node %s, strategy %s",
                    str(out_node), str(out_strtg_pool.get_strtg(out_strtg_idx)))
                continue

    def get_strtg_pool(self) -> ClusterStrategyPool:
        assert self.strategy_pool
        return self.strategy_pool

    def __str__(self) -> str:
        cluster_info = "\ncluster id: " + str(self.unique_id) + "\n  nodes: "
        for node in self.nodes.values():
            cluster_info += str(node) + ", "

        cluster_info += "\n  inputs: " + str(self.args)
        cluster_info += "\n  output: " + str(self.output_node)

        cluster_info += "\nstrategies:\n" + str(self.strategy_pool)
        return cluster_info

    def __repr__(self) -> str:
        return self.__str__()


class MetaGraph:

    def __init__(self, ori_struct) -> None:
        self.ori_struct = ori_struct
        self.input_list = []
        self.op_list = []
        self.output_list = []

        self.node_clusters = []

        self.state_io_map = {}

    def add_input(self, placeholder: MetaNode) -> None:
        self.input_list.append(placeholder)

    def add_node(self, meta_node: MetaNode) -> None:
        self.op_list.append(meta_node)

    def add_output(self, outvars: MetaVar) -> None:
        self.output_list.append(outvars)

    def __str__(self) -> str:
        return_str = f"=====================\n[MetaIR]\n\ninput_list: {self.input_list.__str__()}\n\n"
        for op in self.op_list:
            return_str += f"{op.outvars} <--- [{op.op_name}] --- {op.invars}\n"
        return_str += f"\noutput_list: {self.output_list.__str__()}\n=====================\n"
        return_str += f"\nnode clusters:"
        for cluster in self.node_clusters:
            return_str += str(cluster)
        return_str += "\n=====================\n"
        return return_str

    def __repr__(self) -> str:
        return self.__str__()

    def liveness(self, reserve_input=False):
        liveness_set = set([var.name for var in self.output_list])

        liveness_list = []

        for op in reversed(self.op_list):
            for var in op.invars:
                liveness_set.add(var.name)
            for var in op.outvars:
                if var:  # skip if an output is dangling
                    liveness_set.add(var.name)

            liveness_set_line = copy.deepcopy(liveness_set)
            if reserve_input:
                liveness_set_line.union(set([var.name for var in self.input_list]))

            liveness_list.insert(0, liveness_set_line)

            for var in op.outvars:
                if var:  # skip if an output is dangling
                    liveness_set.remove(var.name)

        return liveness_list

    def build_fine_grain_clusters(self):
        cluster_id = 0
        for node in self.op_list:
            cluster = MetaNodeCluster(unique_id=cluster_id)
            cluster.add_node(node)
            cluster.finalize()
            self.node_clusters.append(cluster)
            cluster_id += 1

    def find_cone_roots(self):
        cone_roots = []
        logger.debug("node num: %s" % len(self.op_list))
        for node in self.op_list:
            down_node_num = 0
            valid_outvars = []
            for out_var in node.outvars:
                if out_var:
                    for down_nd in out_var.down_nodes:
                        if down_nd:
                            valid_outvars.append(out_var)
                            down_node_num += 1

            if down_node_num > 1 or down_node_num == 0:
                cone_roots.append(node)
                continue

            up_node_num = 0
            valid_invars = []
            for in_var in node.invars:
                if in_var and in_var.up_node:
                    valid_invars.append(in_var)
                    up_node_num += 1

            if up_node_num > 1:
                cone_roots.append(node)
                continue
            elif up_node_num == 0:
                continue

            # single up node and single down node
            # Lansong(TODO) we should allow multi-input node as inner node of cone
            assert len(valid_outvars) == 1 and len(valid_invars) == 1
            out_size = valid_outvars[0].get_var_size()
            in_size = valid_invars[0].get_var_size()

            if out_size < in_size:
                cone_roots.append(node)
                continue

        return cone_roots

    def build_cone_cluster(self, nd: MetaNode, root_ids, cluster: MetaNodeCluster):
        cluster.add_node(nd)
        for in_var in nd.invars:
            if in_var and in_var.up_node:
                if in_var.up_node.unique_id not in root_ids:
                    #print("recursive build cone from %s" % str(in_var.up_node))
                    self.build_cone_cluster(in_var.up_node, root_ids, cluster)

    def build_cone_clusters(self):
        cone_roots = self.find_cone_roots()
        logger.debug("root num: %d" % len(cone_roots))
        logger.debug(cone_roots)

        root_ids = set()
        for cone_root in cone_roots:
            root_ids.add(cone_root.unique_id)

        cluster_id = 0
        for cone_root in cone_roots:
            cluster = MetaNodeCluster(unique_id=cluster_id)
            self.build_cone_cluster(cone_root, root_ids, cluster)
            cluster.finalize()
            self.node_clusters.append(cluster)
            cluster_id += 1

    def coarsen(self, coarsen_level: int):
        if coarsen_level == 0:
            self.build_fine_grain_clusters()
        elif coarsen_level == 1:
            self.build_cone_clusters()
        else:
            print("Lansong 2do: support more aggressive coarsening")
            # call cone cluster building instead
            self.build_cone_clusters()

        if mdconfig.log_level <= logging.DEBUG:
            for cluster in self.node_clusters:
                print(str(cluster))

    def dump(self):
        if mdconfig.dump_dir is not None:
            os.makedirs(mdconfig.dump_dir, exist_ok=True)
            filename = os.path.join(mdconfig.dump_dir, "metair.txt")
            with open(filename, "w") as f:
                f.write(self.__str__())
            logger.info(f"MetaIR dump into {filename}")

    def get_input_strategy(self, opt_strategy):

        partial_strategy = {}
        for op in self.op_list:
            if op.is_placeholder:
                if op.unique_key() in opt_strategy:
                    strategy = opt_strategy[op.unique_key()]['strategy'].out_strtg_group[0]
                    partial_strategy[op.outvars[0]] = strategy
                else:
                    logger.warning(f"{op.unique_key()} not found in opt_strategy. (maybe scalar tensor)")

        partial_strategy_list = []

        for var in self.input_list:
            if var in partial_strategy:
                partial_strategy_list.append(partial_strategy[var])
            else:
                partial_strategy_list.append(
                    [SPMD(SPMD.REPLICATE),
                     SPMD(SPMD.REPLICATE)])

        return partial_strategy_list


easydist/metashard/metair.py
alibaba/easydist
# Copyright (c) 2023, Alibaba Group;
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================

import copy
import logging

import numpy as np

from easydist import platform
from easydist.metashard.halo import HaloInfo, halo_padding
from easydist.metashard.annotation import ShardDim, ShardAnnotation
from easydist.metashard.combination import try_combination, HaloHint

logger = logging.getLogger(__name__)


def get_shard_size(dim_to_shard, shard_size):
    gcd = np.gcd.reduce(dim_to_shard)
    if gcd < shard_size:
        return gcd
    # elif gcd % shard_size == 0:
    #     return shard_size
    # return -1
    return shard_size


def get_shard_data(input_args, shard_size, dim_idx, haloinfo, chunk):
    chunk_data = platform.chunk(platform.clone(input_args), chunks=chunk, dim=dim_idx)
    chunk_shard_data = [
        platform.chunk(chunk_, chunks=shard_size, dim=dim_idx) for chunk_ in chunk_data
    ]
    shard_data = [
        platform.concatenate([t[i] for t in chunk_shard_data], dim_idx) for i in range(shard_size)
    ]

    return halo_padding(shard_data, haloinfo=haloinfo)


def check_prompt(flat_input_args, prompt_annotation):
    tensor_args = [i for i in flat_input_args if isinstance(flat_input_args, platform.Tensor)]
    if len(tensor_args) == len(prompt_annotation):
        for idx in range(len(tensor_args)):
            if tensor_args.ndim != len(prompt_annotation[idx]):
                return False
        return True
    return False


class MetaOp:

    def __init__(self, func, input_args, shard_size=2, name=None) -> None:
        self.func = func
        self.name = name
        if self.name is None:
            self.name = self.func.__name__
        self.input_args = input_args
        self.shard_size = shard_size
        self.flat_input_args, self.input_pytree_spec = platform.tree_flatten(input_args)

    def exec_platform(self, flat_input_args):
        if platform.get_backend() in ["torch", "tvm"]:
            args, kwargs = platform.tree_unflatten(flat_input_args, self.input_pytree_spec)
            return self.func(*args, **kwargs)
        elif platform.get_backend() == "jax":
            # (NOTE) disable_jit here for jax >= 4.7, the pjit primitive doesn't accept chunked input
            import jax
            with jax.disable_jit():
                flat_input_args = platform.tree_unflatten(flat_input_args, self.input_pytree_spec)
                if len(flat_input_args) == 3 and isinstance(flat_input_args[2], dict):
                    subfuns, invars, bind_params = flat_input_args
                    return self.func(*subfuns, *invars, **bind_params)
                return self.func(*flat_input_args)

    def exec(self, shard_annotation=None, priority_shard_dim_id=None, haloinfo=None):
        local_flat_input_args = self.flat_input_args

        tensor_list = [i for i in local_flat_input_args if isinstance(i, platform.Tensor)]

        if shard_annotation:
            dim_to_shard = []
            for tensor_ann_idx, tensor_ann in enumerate(shard_annotation):
                for dim_idx, dim in enumerate(tensor_ann):
                    if dim.shard_dim_id == priority_shard_dim_id:
                        dim_to_shard.append(tensor_list[tensor_ann_idx].shape[dim_idx])
                        break
            if len(dim_to_shard) == 0:
                raise RuntimeError(f"can not find priority_shard_dim_id {priority_shard_dim_id}")
            shard_size = get_shard_size(dim_to_shard, self.shard_size)
            if shard_size < self.shard_size:
                raise RuntimeError("can not find suitable shard_size")
            else:
                sharded_out = []
                for shard_idx in range(shard_size):
                    sharded_flat_input_args = []
                    tensor_arg = 0
                    for input_args in local_flat_input_args:
                        if isinstance(input_args, platform.Tensor):
                            tensor_ann = shard_annotation[tensor_arg]
                            no_shard_here = True
                            for dim_idx, dim in enumerate(tensor_ann):
                                if dim.shard_dim_id == priority_shard_dim_id:
                                    shard_data = get_shard_data(input_args, shard_size, dim_idx,
                                                                haloinfo, dim.chunk)[shard_idx]
                                    sharded_flat_input_args.append(shard_data)
                                    no_shard_here = False
                                    break
                            if no_shard_here:
                                sharded_flat_input_args.append(platform.clone(input_args))
                            tensor_arg += 1
                        else:
                            sharded_flat_input_args.append(input_args)

                    sharded_out.append(self.exec_platform(sharded_flat_input_args))
                return sharded_out

        else:
            return self.exec_platform(local_flat_input_args)

    def _try_sharding(self, fixed_annotation, subsequence_annotation, global_output, start_dim=0):
        shard_dim_id_flag = self.__shard_dim_id_flag

        if self.__find_new_strategy:
            return

        if len(subsequence_annotation) == 0:
            try:
                sharded_output = self.exec(shard_annotation=fixed_annotation,
                                           priority_shard_dim_id=shard_dim_id_flag)
            except RuntimeError as e:
                logger.debug(f"[{fixed_annotation}] {e}")
                return
            except:
                logger.debug(f"[{fixed_annotation}] run op.exec failed")
                return

            haloinfo = None
            combination_func = try_combination(sharded_output, global_output)
            if isinstance(combination_func, HaloHint):
                ecpt = combination_func
                if ecpt.halo <= 0:
                    ecpt.halo = 1
                if isinstance(sharded_output[0], tuple) or isinstance(sharded_output[0], list):
                    try_max_halo = sharded_output[0][ecpt.out_idx].shape[ecpt.dim] // 2
                else:
                    try_max_halo = sharded_output[0].shape[ecpt.dim] // 2
                for halo in range(ecpt.halo, try_max_halo):
                    sharded_output = self.exec(shard_annotation=fixed_annotation,
                                               priority_shard_dim_id=shard_dim_id_flag,
                                               haloinfo=HaloInfo(halo, ecpt.dim))
                    combination_func = try_combination(sharded_output, global_output)
                    if isinstance(combination_func, HaloHint):
                        combination_func = None
                    if combination_func:
                        haloinfo = HaloInfo(halo, ecpt.dim)
                        break

            if combination_func is not None and not isinstance(combination_func, HaloHint):
                self.__combination_ann[shard_dim_id_flag] = combination_func
                # inject haloinfo
                fixed_annotation.inject_haloinfo(haloinfo, shard_dim_id_flag)

                self.__sharding_annotion = copy.deepcopy(fixed_annotation)
                self.__find_new_strategy = True
            else:
                logger.debug(f"[{fixed_annotation}] combination failed")
            return

        for dim in range(start_dim, len(subsequence_annotation[0])):
            if subsequence_annotation[0][dim].shard_dim_id != 0:
                continue
            try_annotation = copy.deepcopy(subsequence_annotation[0])
            try_annotation[dim] = ShardDim.get_shard_dim(shard_dim_id_flag)
            self._try_sharding(fixed_annotation + ShardAnnotation([try_annotation]),
                               subsequence_annotation[1:], global_output)

        self._try_sharding(fixed_annotation + ShardAnnotation([subsequence_annotation[0]]),
                           subsequence_annotation[1:], global_output)

    def sharding_discovery(self, prompt_annotation=None):

        self.__combination_ann = {}
        # validate the prompt_annotation
        if prompt_annotation is not None and check_prompt(self.flat_input_args, prompt_annotation):
            max_shard_dim_id = prompt_annotation.get_max_shard_dim_id()
            for shard_dim_id in range(1, max_shard_dim_id + 1):
                combination_func = self.sharding_discovery_with_preset(
                    prompt_annotation, shard_dim_id)

                if combination_func is not None:
                    self.__combination_ann[shard_dim_id] = combination_func
                else:
                    break

        validate_max_shard_dim_id = len(self.__combination_ann)

        if validate_max_shard_dim_id >= 1:
            self.__sharding_annotion = prompt_annotation.clear_shard_dim(validate_max_shard_dim_id)
        else:
            # init the annotation with all NoShardDim
            self.__sharding_annotion = ShardAnnotation.init_from_input_args(self.flat_input_args)

        global_output = self.exec()

        self.__shard_dim_id_flag = validate_max_shard_dim_id + 1
        self.__find_new_strategy = False

        if len(self.__sharding_annotion) == 0:
            return self.__sharding_annotion, self.__combination_ann

        fixed_annotation = ShardAnnotation([])
        subsequence_annotation = self.__sharding_annotion
        start_dim = 0
        while True:
            self._try_sharding(fixed_annotation, subsequence_annotation, global_output, start_dim)
            if self.__find_new_strategy:

                # if shard_dim_id_flag not in sharding_annotion, then break
                max_shard_dim_id_now = 0
                for ann in self.__sharding_annotion:
                    for dim in ann:
                        max_shard_dim_id_now = max(max_shard_dim_id_now, dim.shard_dim_id)
                if max_shard_dim_id_now != self.__shard_dim_id_flag:
                    break

                self.__find_new_strategy = False
                # find the first dim for shard_dim_id_flag now
                # and update fixed_annotation, subsequence_annotation, start_dim
                fixed_annotation = ShardAnnotation([])
                finded = False
                for ann_idx, ann in enumerate(self.__sharding_annotion):
                    for idx, dim in enumerate(ann):
                        if self.__shard_dim_id_flag == dim.shard_dim_id:
                            finded = True
                            start_dim = idx + 1
                            break
                    if finded:
                        break
                    fixed_annotation += ShardAnnotation([ann])
                subsequence_annotation = self.__sharding_annotion[ann_idx:]
                if ann_idx == len(self.__sharding_annotion) - 1 and start_dim == len(
                        self.__sharding_annotion[-1]):
                    break
                self.__shard_dim_id_flag += 1
            else:
                break

        logger.debug(f"sharding_annotion of {self.name}: {self.__sharding_annotion}")

        return self.__sharding_annotion, self.__combination_ann

    def sharding_discovery_with_preset(self, sharding_annotion, priority_shard_dim_id=1):
        global_output = self.exec()

        try:
            sharded_output = self.exec(shard_annotation=sharding_annotion,
                                       priority_shard_dim_id=priority_shard_dim_id)
        except RuntimeError as e:
            logger.debug(f"[{sharding_annotion}] {e}")
            return
        except:
            logger.debug(f"[{sharding_annotion}] run op.exec failed")
            return

        combination_func = try_combination(sharded_output, global_output)
        if combination_func is not None and not isinstance(combination_func, HaloHint):
            return combination_func

easydist/metashard/metaop.py
alibaba/easydist
# Copyright (c) 2023, Alibaba Group;
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================

import functools
import operator

from easydist.metashard.annotation import ShardAnnotation, ShardDim
from easydist.metashard.combination import CombinationFunc

EXTEND_VIEW = False


def get_next_non_one(shape_, idx_):
    if idx_ < len(shape_):
        while shape_[idx_] == 1:
            idx_ += 1
            if idx_ >= len(shape_):
                break
    return idx_


def view_propagation(input_shape, output_shape, world_size=1):

    if -1 in output_shape:
        numel = functools.reduce(operator.mul, input_shape)
        dim_size = -1 * numel // functools.reduce(operator.mul, output_shape)
        output_shape[output_shape.index(-1)] = dim_size

    sharding_ann = ShardAnnotation([[ShardDim(0)] * len(input_shape)])
    combination_ann = {}

    input_idx = get_next_non_one(input_shape, 0)
    output_idx = get_next_non_one(output_shape, 0)

    shard_dim = 1

    while input_idx < len(input_shape):
        if input_shape[input_idx] == output_shape[output_idx]:
            # [**, A, **] -> [**, A, **]
            if input_shape[input_idx] >= world_size:
                sharding_ann[0][input_idx] = ShardDim(shard_dim)
                combination_ann[shard_dim] = functools.partial(CombinationFunc.gather,
                                                               dim=output_idx)
                shard_dim += 1
            input_idx = get_next_non_one(input_shape, input_idx + 1)
            output_idx = get_next_non_one(output_shape, output_idx + 1)
        elif input_shape[input_idx] > output_shape[output_idx]:
            # [**, A, **] -> [**, a1, a2, **]
            leftmost_idx = output_idx
            accum_shape_ = output_shape[output_idx]
            for o_idx in range(output_idx + 1, len(output_shape)):
                accum_shape_ *= output_shape[o_idx]
                if accum_shape_ == input_shape[input_idx]:
                    if output_shape[leftmost_idx] >= world_size:
                        sharding_ann[0][input_idx] = ShardDim(shard_dim)
                        combination_ann[shard_dim] = functools.partial(CombinationFunc.gather,
                                                                       dim=leftmost_idx)
                        shard_dim += 1
                    output_idx = get_next_non_one(output_shape, o_idx + 1)
                    input_idx = get_next_non_one(input_shape, input_idx + 1)
                    break
        else:
            # [**, a1, a2, **] -> [**, A, **]
            leftmost_idx = input_idx
            accum_shape_ = input_shape[input_idx]
            for i_idx in range(input_idx + 1, len(input_shape)):
                accum_shape_ *= input_shape[i_idx]
                if accum_shape_ == output_shape[output_idx]:
                    if EXTEND_VIEW:
                        chunk_size_ = 1
                        for sub_idx in range(input_idx, i_idx + 1):
                            sharding_ann[0][sub_idx] = ShardDim(shard_dim)
                            combination_ann[shard_dim] = functools.partial(CombinationFunc.gather,
                                                                           dim=output_idx,
                                                                           chunk=chunk_size_)
                            chunk_size_ *= input_shape[sub_idx]
                            shard_dim += 1
                    else:
                        if input_shape[input_idx] >= world_size:
                            sharding_ann[0][input_idx] = ShardDim(shard_dim)
                            combination_ann[shard_dim] = functools.partial(CombinationFunc.gather,
                                                                           dim=output_idx)
                            shard_dim += 1

                    output_idx = get_next_non_one(output_shape, output_idx + 1)
                    input_idx = get_next_non_one(input_shape, i_idx + 1)
                    break

    return {'sharding_ann': sharding_ann, 'combination_ann': combination_ann}


def view_propagation_preset(input_shape, output_shape, preset_anno):
    accum_size = 1
    for idx, ann in enumerate(preset_anno[0]):
        if ann.shard_dim_id != 0:
            break
        accum_size *= input_shape[idx]

    chunk = preset_anno[0][idx].chunk

    out_accum_size = 1
    out_idx = 0
    while out_accum_size < accum_size:
        out_accum_size *= output_shape[out_idx]
        out_idx += 1

    if out_accum_size == accum_size:
        accum_chunk = 1
        if chunk == accum_chunk:
            return functools.partial(CombinationFunc.gather, dim=out_idx)
        for o_idx in range(out_idx, len(output_shape)):
            if chunk == accum_chunk:
                return functools.partial(CombinationFunc.gather, dim=o_idx)
            accum_chunk *= output_shape[o_idx]

easydist/metashard/view_propagation.py
alibaba/easydist
# Copyright (c) 2023, Alibaba Group;
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================

import os
import logging
import importlib

logger = logging.getLogger(__name__)

EASYDIST_BACKEND = None

__all__ = [
    "add", "equal", "zeros_like", "min", "max", "allclose", "concatenate", "chunk", "narrow",
    "Tensor", "tree_flatten", "tree_unflatten", "clone", "from_numpy"
]


def backend_valid(_backend):
    return _backend in {"torch", "jax", "tvm"}


def init_backend(backend="torch"):
    assert backend_valid(backend)
    global EASYDIST_BACKEND
    EASYDIST_BACKEND = backend
    modules = importlib.import_module("." + backend, __name__)
    for val in __all__:
        exec("globals()['%s'] = modules.%s" % (val, val))
    logger.info(f"========= EasyDist init with backend {backend}. =========")


def get_backend():
    global EASYDIST_BACKEND
    return EASYDIST_BACKEND


for val in __all__:
    exec("globals()['%s'] = None" % val)

easydist/platform/__init__.py
alibaba/easydist
# Copyright (c) 2023, Alibaba Group;
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================

from functools import partial

import jax

add = jax.numpy.add
equal = jax.numpy.array_equal
zeros_like = jax.numpy.zeros_like
min = jax.numpy.minimum
max = jax.numpy.maximum
allclose = partial(jax.numpy.allclose, rtol=5e-3, atol=5e-03)


def concatenate(tensors, dim=0):
    return jax.numpy.concatenate(tensors, axis=dim)


def chunk(input, chunks, dim=0):
    return jax.numpy.array_split(input, chunks, axis=dim)


def narrow(input, dim, start, length):
    indices = jax.numpy.asarray(range(start, start + length))
    return jax.numpy.take(input, indices, axis=dim)


Tensor = jax.Array

tree_flatten = jax.tree_util.tree_flatten


def tree_unflatten(values, spec):
    return jax.tree_util.tree_unflatten(spec, values)


clone = jax.numpy.copy

from_numpy = jax.numpy.array

easydist/platform/jax.py
alibaba/easydist
# Copyright (c) 2023, Alibaba Group;
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================

from functools import partial

import torch
import torch.utils._pytree as pytree

add = torch.Tensor.add
equal = torch.equal
zeros_like = torch.zeros_like
min = torch.min
max = torch.max
allclose = partial(torch.allclose, rtol=1e-3, atol=1e-07)
concatenate = torch.concatenate
chunk = torch.chunk
narrow = torch.narrow

Tensor = torch.Tensor

tree_flatten = pytree.tree_flatten
tree_unflatten = pytree.tree_unflatten

clone = torch.clone
from_numpy = torch.from_numpy

easydist/platform/torch.py
alibaba/easydist
# Copyright (c) 2023, Alibaba Group;
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================

import tvm
import numpy
import torch.utils._pytree as pytree


def add(a, b):
    return tvm.nd.array(numpy.add(a.numpy(), b.numpy()))


def equal(a, b):
    return numpy.array_equal(a.numpy(), b.numpy())


def zeros_like(input_):
    return tvm.nd.array(numpy.zeros_like(input_.numpy()))


def min(a, b):
    return tvm.nd.array(numpy.minimum(a.numpy(), b.numpy()))


def max(a, b):
    return tvm.nd.array(numpy.maximum(a.numpy(), b.numpy()))


def allclose(a, b):
    return numpy.allclose(a.numpy(), b.numpy(), rtol=5e-3, atol=5e-03)


def concatenate(tensors, dim=0):
    return tvm.nd.array(numpy.concatenate([t.numpy() for t in tensors], axis=dim))


def chunk(input, chunks, dim=0):
    return [tvm.nd.array(i) for i in numpy.array_split(input.numpy(), chunks, axis=dim)]


def narrow(input, dim, start, length):
    indices = numpy.asarray(range(start, start + length))
    return tvm.nd.array(numpy.take(input.numpy(), indices, axis=dim))


Tensor = tvm.nd.NDArray

tree_flatten = pytree.tree_flatten
tree_unflatten = pytree.tree_unflatten


def clone(input_):
    return tvm.nd.array(numpy.copy(input_.numpy()))


from_numpy = tvm.nd.array

easydist/platform/tvm.py
alibaba/easydist
# Copyright (c) 2023, Alibaba Group;
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================

import os

import torch
from torch._functorch import config
from torch.distributed._tensor.op_schema import OpSchema

from .api import compile, enable_transform, get_input_strategy, easydist_shard
from .bridge import shard_module, torch2meta_graph
from .passes.sharding import sharding_transform
from .device_mesh import set_device_mesh, get_device_mesh
from .sharding_interpreter import EDTorchShardingAnn
from .spmd_prop_rule import *

# disable with torch <= 2.0.1
if hasattr(config, "use_fake_tensor"):
    config.use_fake_tensor = False

__all__ = [
    'EDTorchShardingAnn', 'sharding_transform', 'set_device_mesh', 'get_device_mesh',
    'torch2meta_graph', 'compile', 'easydist_shard', 'enable_transform', 'shard_module',
    'get_input_strategy'
]


def easydist_setup_torch(device, allow_tf32):
    torch.backends.cuda.matmul.allow_tf32 = allow_tf32
    torch.backends.cudnn.allow_tf32 = allow_tf32

    # this environment used for cuda graph
    os.environ["NCCL_ASYNC_ERROR_HANDLING"] = "0"


# hotfix for PyTorch 2.1.0
if "2.1.0" in torch.__version__:
    def hot_fix_opschema_hash(self):
        # NOTE: we turn kwargs_schema into a frozenset to hash as it would not be nested dict
        frozen_set_kwargs_schema = frozenset(self.kwargs_schema.items())
        return hash(
            (
                self.func_schema.name,
                tuple(self.func_schema.arguments),
                tuple(self.func_schema.returns),
                self.func_schema.name,
                tuple(tuple(e) if isinstance(e, list) else e for e in self.args_schema),
                frozen_set_kwargs_schema,
            )
        )

    OpSchema.__hash__ = hot_fix_opschema_hash

easydist/torch/__init__.py
alibaba/easydist
# Copyright (c) 2023, Alibaba Group;
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================

import time
import logging

import rich
import torch
from torch._functorch.aot_autograd import aot_function, make_boxed_compiler

from easydist.metashard import metair
from easydist.autoflow.solver import AutoFlowSolver
import easydist.config as mdconfig
from .sharding_interpreter import EDTorchShardingAnn
from .utils import _sharding_ann_env
from .bridge import torch2meta_graph, get_torch_sharding_strategy
from .passes import eliminate_detach, fix_addmm_bias, fix_convoluation_bias, fix_embedding, sharding_transform
from .passes.sharding import get_device_mesh

logger = logging.getLogger(__name__)

ENABLE_TRANSFORM = False
BW_CONSTRAINTS = None
INPUT_STRATEGY = None


def get_input_strategy():
    global INPUT_STRATEGY
    return INPUT_STRATEGY


def enable_transform():
    global ENABLE_TRANSFORM
    ENABLE_TRANSFORM = True


def _get_output_strategy(opt_strategy, meta_graph, input_strategy):
    partial_strategy = {}
    for op in meta_graph.op_list:
        op_key = op.unique_key()
        if op_key in opt_strategy:
            for idx, var in enumerate(op.outvars):
                if var in meta_graph.output_list:
                    strategy = opt_strategy[op_key]['strategy'].out_strtg_group.get_var_strtg(idx)
                    partial_strategy[var] = strategy

    for var in meta_graph.input_list:
        if var in meta_graph.output_list:
            if var in input_strategy:
                partial_strategy[var] = input_strategy[var]

    return partial_strategy


def _get_input_strategy(opt_strategy, meta_graph):
    partial_strategy = {}
    for op in reversed(meta_graph.op_list):
        op_key = op.unique_key()
        if op_key in opt_strategy:
            for idx, var in enumerate(op.invars):
                if var in meta_graph.input_list:
                    strategy = opt_strategy[op_key]['strategy'].in_strtg_group.get_var_strtg(idx)
                    partial_strategy[var] = strategy

    partial_strategy_list = []

    for var in meta_graph.input_list:
        if var in partial_strategy:
            partial_strategy_list.append(partial_strategy[var])
        else:
            partial_strategy_list.append(
                [metair.SPMD(metair.SPMD.REPLICATE),
                 metair.SPMD(metair.SPMD.REPLICATE)])

    return partial_strategy, partial_strategy_list


@make_boxed_compiler
def easydist_shard(fx_module: torch.fx.GraphModule, inps):

    global BW_CONSTRAINTS, INPUT_STRATEGY

    fx_module = fix_embedding(fx_module)
    fx_module = fix_addmm_bias(fx_module)
    fx_module = fix_convoluation_bias(fx_module)
    fx_module = eliminate_detach(fx_module)
    fx_module.recompile()
    if mdconfig.log_level <= logging.DEBUG:
        print(fx_module.graph)

    with _sharding_ann_env():
        start_t = time.perf_counter()
        sharding_interpreter = EDTorchShardingAnn(fx_module)
        sharding_info, fwd_shape_info = sharding_interpreter.run(*inps)
        logger.info(f"[EDTorchShardingAnn.time]:\t {time.perf_counter() - start_t} s.")
        if mdconfig.log_level <= logging.DEBUG:
            rich.print("sharding_info:\n", sharding_info)
            rich.print("fwd_shape_info:\n", fwd_shape_info)

    meta_graph = torch2meta_graph(fx_module, sharding_info, fwd_shape_info)

    if mdconfig.log_level <= logging.DEBUG:
        rich.print(meta_graph)

    device_mesh = get_device_mesh()
    device_mesh_shape = (device_mesh.size(0), device_mesh.size(1))

    solver = AutoFlowSolver(device_mesh_shape, BW_CONSTRAINTS)
    solver.add_graph(meta_graph)
    start_t = time.perf_counter()
    count_invars = BW_CONSTRAINTS is None
    opt_strategy = solver.ilp_optimize(count_invars)
    logger.info(f"[AutoFlowSolver.time]:\t {time.perf_counter() - start_t} s.")
    # start_t = time.perf_counter()
    # beam_search_strategy = solver.beam_search()
    # logger.info(f"[AutoFlowSolver.beam_search]: {time.perf_counter() - start_t} s.")

    if mdconfig.log_level <= logging.DEBUG:
        rich.print(opt_strategy)

    if BW_CONSTRAINTS is None:
        strategy_map, INPUT_STRATEGY = _get_input_strategy(opt_strategy, meta_graph)
        BW_CONSTRAINTS = _get_output_strategy(opt_strategy, meta_graph, strategy_map)

    if ENABLE_TRANSFORM:
        sharding_strategy = get_torch_sharding_strategy(fx_module, opt_strategy)

        if mdconfig.log_level <= logging.DEBUG:
            print(sharding_strategy)

        fx_module = sharding_transform(fx_module, sharding_strategy)
        fx_module = fix_embedding(fx_module, recover=True)

        if mdconfig.log_level <= logging.DEBUG:
            print(fx_module.graph)

    return fx_module


def md_aot_module(mod: torch.nn.Module, *args, **kwargs) -> torch.nn.Module:
    """
    Traces the forward and backward graph of :attr:`mod` using torch dispatch
    tracing mechanism. It is wrapper function, that underneath uses
    :func:`aot_function` to perform tracing and compilation.

    :func:`aot_module` lifts the parameters and buffers of ``nn.Module`` as inputs
    to a new callable which is then compiled through :func:`aot_function`.

    .. warning::
        This API is experimental and likely to change.

    Args:
        mod (Callable): A ``nn.Module`` module.
        args : args to be passed to :func:`aot_function`
        kwargs : kwargs to be passed to :func:`aot_function`

    Returns:
        Returns a ``nn.Module`` that retains the eager behavior of the original
        :attr:`mod`, but with forward and backward graph compiled.

    """
    # See Note: [Fake Modules and AOTAutograd]
    torch._dynamo.utils.assert_no_fake_params_or_buffers(mod)

    def functional_call(named_params, named_buffers, *args, **kwargs):
        params_and_buffers = {**named_params, **named_buffers}
        return torch.func.functional_call(mod, params_and_buffers, args, kwargs)

    named_params = dict(mod.named_parameters(remove_duplicate=False))
    named_buffers = dict(mod.named_buffers(remove_duplicate=False))
    num_params_buffers = len(named_params) + len(named_buffers)
    compiled_f = aot_function(functional_call,
                              num_params_buffers=num_params_buffers,
                              *args,
                              **kwargs)

    class AOTModule(torch.nn.Module):

        def __init__(self):
            super().__init__()
            self.orig_module = mod

        def forward(self, *args, **kwargs):
            named_params = dict(mod.named_parameters(remove_duplicate=False))
            named_buffers = dict(mod.named_buffers(remove_duplicate=False))
            return compiled_f(
                named_params,
                named_buffers,
                *args,
                **kwargs,
            )

    return AOTModule()


def compile(mod: torch.nn.Module):
    return md_aot_module(mod, fw_compiler=easydist_shard, bw_compiler=easydist_shard)

easydist/torch/api.py
alibaba/easydist
# Copyright (c) 2023, Alibaba Group;
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================

import os
import torch
from torch.fx.node import Node, _get_qualified_name
import torch.distributed.distributed_c10d as c10d

import torch.distributed._tensor as spmd
from torch.distributed._tensor import distribute_tensor
import torch.utils._pytree as pytree

from easydist.metashard.combination import ReduceOp
from easydist.metashard import metair
from easydist.metashard.metair import MetaGraph, MetaNode, MetaVar
from easydist.utils import rsetattr, rgetattr
import easydist.config as mdconfig

from .passes.sharding import get_device_mesh

ABSTRACT_DTYPE = {
    torch.float64: "float64",
    torch.float32: "float32",
    torch.float16: "float16",
    torch.int32: "int32",
    torch.int64: "int64",
    torch.bool: "bool",
    torch.uint8: "uint8",
    torch.complex64: "complex64",
}


def to_torch_spmd(meta_spmd):
    if meta_spmd.state == metair.SPMD.SHARD:
        return spmd.Shard(dim=meta_spmd.args["dim"])
    elif meta_spmd.state == metair.SPMD.PARTIAL:
        mapping_ops = {
            ReduceOp.SUM: c10d.ReduceOp.RedOpType.SUM,
            ReduceOp.MAX: c10d.ReduceOp.RedOpType.MAX,
            ReduceOp.MIN: c10d.ReduceOp.RedOpType.MIN,
        }
        return spmd.placement_types._Partial(reduce_op=mapping_ops[meta_spmd.args["ops"]])
    elif meta_spmd.state == metair.SPMD.REPLICATE:
        return spmd.Replicate()


def materialize(x, device):
    if isinstance(x, torch.Tensor) and x.is_meta:
        if x.dtype == torch.bool:
            return torch.rand(x.size(), dtype=torch.float, device=device) > 0.5
        elif torch.is_floating_point(x):
            return torch.rand(x.size(), dtype=x.dtype, device=device)
        else:
            return torch.randint(high=8, size=x.size(), dtype=x.dtype, device=device)
    return x


def shard_module(model, input_, input_strategy, device="cuda"):
    mesh = get_device_mesh()

    input_strategy = [[to_torch_spmd(i) for i in var_strategy] for var_strategy in input_strategy]

    idx = 0
    for name in dict(model.named_parameters()):

        tensor_data = rgetattr(model, name).data
        tensor_data = materialize(tensor_data, device=device)

        rsetattr(
            model, name,
            torch.nn.parameter.Parameter(distribute_tensor(tensor_data, mesh,
                                                           input_strategy[idx])))

        with torch.no_grad():
            rsetattr(model, name + ".grad", torch.empty_like(rgetattr(model, name).data))
        idx += 1

    for name in dict(model.named_buffers()):

        tensor_data = rgetattr(model, name).data
        tensor_data = materialize(tensor_data, device=device)

        rsetattr(model, name, distribute_tensor(tensor_data, mesh, input_strategy[idx]))
        idx += 1

    shard_input = []
    for tensor in input_:
        tensor = materialize(tensor, device=device)
        shard_input.append(distribute_tensor(tensor, mesh, input_strategy[idx]))
        idx += 1

    return shard_input


def torch2meta_graph(fx_module: torch.fx.GraphModule, state_tensor_num, sharding_info,
                     meta_info) -> MetaGraph:
    meta_graph = MetaGraph(fx_module)
    meta_node_map = {}
    meta_var_map = {}
    MetaNode.clear_id_counter()
    MetaVar.clear_id_counter()
    output_names = []
    # 1. create MetaVar for each output of fx graph node except list or tuple
    #    and create MetaNode for each fx graph node except getitem

    for node in fx_module.graph.nodes:
        if node.op == "call_function":
            # 1.1. create MetaVar
            outvars = []
            compact_out_idx_tbl = []
            compact_out_idx = 0
            # NOTE: list and tuple are followed by getitem
            if isinstance(meta_info[node.name], list) or isinstance(meta_info[node.name], tuple):
                for idx, var_meta in enumerate(meta_info[node.name]):
                    if var_meta is not None and var_meta != {}:
                        compact_out_idx_tbl.append(compact_out_idx)
                        compact_out_idx = compact_out_idx + 1
                    else:
                        compact_out_idx_tbl.append(-1)

                # don't create MetaVar. Will fill it with getitem's MetaVar
                if compact_out_idx > 0:
                    outvars = [None] * compact_out_idx
            else:
                compact_out_idx_tbl = [0]
                meta_var = MetaVar(name=node.name,
                                   shape=meta_info[node.name]["shape"],
                                   dtype=ABSTRACT_DTYPE[meta_info[node.name]["dtype"]])

                meta_var_map[node.name] = meta_var
                outvars.append(meta_var)

            op_name = _get_qualified_name(node.target)
            if op_name == "_operator.getitem":
                # not create MetaNode
                continue

            # 1.2. create MetaNode
            node_sharding_info = None
            if op_name in sharding_info:
                def _gen_meta(arg):
                    if isinstance(arg, Node):
                        return torch.empty(meta_info[arg.name]["shape"],
                                           dtype=meta_info[arg.name]["dtype"],
                                           device="meta")
                    else:
                        # primitive data type: int, float, etc
                        return arg

                args_meta = pytree.tree_map(_gen_meta, node.args)
                args_meta = str(tuple(args_meta)) + ' | ' + str(node.kwargs)
                if args_meta in sharding_info[op_name]:
                    node_sharding_info = sharding_info[op_name][args_meta]

            node_args_flatten = pytree.tree_flatten(node.args)[0]
            compact_in_idx_tbl = []
            compact_in_idx = 0
            for arg in node_args_flatten:
                if isinstance(arg, Node):
                    compact_in_idx_tbl.append(compact_in_idx)
                    compact_in_idx = compact_in_idx + 1
                else:
                    compact_in_idx_tbl.append(-1)
            invars = [None] * compact_in_idx

            meta_node = MetaNode(name=node.name,
                                 op_name=op_name,
                                 invars=invars,
                                 outvars=outvars,
                                 sharding_info=node_sharding_info)
            meta_node.compact_out_idx_tbl = compact_out_idx_tbl
            meta_node.compact_in_idx_tbl = compact_in_idx_tbl
            meta_node_map[node.name] = meta_node
            meta_graph.add_node(meta_node)
        elif node.op in ["placeholder", "get_attr"]:
            if meta_info[node.name] != {}:
                # 1.1. create MetaVar
                meta_var = MetaVar(name=node.name,
                                   shape=meta_info[node.name]["shape"],
                                   dtype=ABSTRACT_DTYPE[meta_info[node.name]["dtype"]])

                meta_var_map[node.name] = meta_var

                # 1.2. create MetaNode
                node_sharding_info = None
                if node.op in sharding_info:
                    arg_meta_tensor = torch.empty(meta_info[node.name]["shape"],
                                                  dtype=meta_info[node.name]["dtype"],
                                                  device="meta")
                    args_meta = str(arg_meta_tensor)
                    if args_meta in sharding_info[node.op]:
                        node_sharding_info = sharding_info[node.op][args_meta]

                meta_node = MetaNode(name=node.name,
                                     op_name=node.op,
                                     invars=[],
                                     outvars=[meta_var],
                                     sharding_info=node_sharding_info,
                                     is_placeholder=True)
                meta_node.compact_out_idx_tbl = [0]
                meta_node.compact_in_idx_tbl = []
                meta_node_map[node.name] = meta_node
                meta_graph.add_node(meta_node)

                meta_graph.add_input(meta_node)
        elif node.op == "output":
            output_names = [arg.name for arg in node.args[0] if arg is not None]

    # 2. update connection between MetaNode and MetaVar
    for node in fx_module.graph.nodes:
        if node.op == "call_function":
            op_name = _get_qualified_name(node.target)
            if op_name == "_operator.getitem":
                # connect MetaVar of this node to MetaNode of its up node
                arg, idx = node.args
                assert isinstance(arg, Node)
                assert isinstance(idx, int)
                arg_meta_node = meta_node_map[arg.name]  # find meta node of getitem's up node
                meta_var = meta_var_map[node.name]  # find getitem's output var
                arg_meta_node.set_out_var(meta_var, idx)
                continue

            node_args_flatten = pytree.tree_flatten(node.args)[0]
            for idx, arg in enumerate(node_args_flatten):
                if isinstance(arg, Node):
                    meta_node = meta_node_map[node.name]
                    in_var = meta_var_map[arg.name]
                    meta_node.set_in_var(in_var, idx)

    for ouput_n in output_names:
        meta_graph.add_output(meta_var_map[ouput_n])

    # add state io map for input/output sharding mismatch cost
    state_io_map = {}
    for i in range(state_tensor_num):
        state_io_map[meta_graph.input_list[i]] = meta_graph.output_list[i]
    meta_graph.state_io_map = state_io_map

    meta_graph.coarsen(coarsen_level=mdconfig.coarsen_level)

    return meta_graph


def get_torch_sharding_strategy(fx_module: torch.fx.GraphModule, opt_strategy):
    sharding_strategy = {}

    for node in fx_module.graph.nodes:
        if node.op == "call_function":
            op_name = _get_qualified_name(node.target)
            if op_name != "_operator.getitem":
                unique_key = node.name
                if unique_key in opt_strategy:
                    sharding_strategy[node.name] = [[
                        to_torch_spmd(ii) for ii in i
                    ] for i in opt_strategy[unique_key]['strategy'].in_strtg_group]
                else:
                    unique_key = f"{op_name}_[{node.name}]"
                    if unique_key in opt_strategy:
                        sharding_strategy[node.name] = [[
                            to_torch_spmd(ii) for ii in i
                        ] for i in opt_strategy[unique_key]['strategy'].in_strtg_group]

    return sharding_strategy

easydist/torch/bridge.py
alibaba/easydist
# Copyright (c) 2023, Alibaba Group;
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================

import logging
import functools
import operator

import numpy
import torch
from torch.distributed._tensor import DeviceMesh

from easydist.metashard import metair
from easydist.utils.testing import TorchMockDeviceMesh

logger = logging.getLogger(__name__)

TORCH_DEVICE_MESH = None


def set_device_mesh(device_mesh):
    global TORCH_DEVICE_MESH
    TORCH_DEVICE_MESH = device_mesh

    if device_mesh.size(0) == 1:
        metair.DEVICE_MESH_1D = 0
    elif device_mesh.size(1) == 1:
        metair.DEVICE_MESH_1D = 1

    logger.info(f"set_device_mesh: {device_mesh}")


def get_device_mesh():
    global TORCH_DEVICE_MESH
    return TORCH_DEVICE_MESH


def device_mesh_world_size(device_mesh=None):
    if device_mesh is None:
        device_mesh = get_device_mesh()

    if device_mesh is None:
        return None

    if isinstance(device_mesh, TorchMockDeviceMesh):
        device_mesh_shape = get_device_mesh().shape
    elif isinstance(device_mesh, DeviceMesh):
        device_mesh_shape = tuple(get_device_mesh().mesh.shape)

    return functools.reduce(operator.mul, device_mesh_shape)


def device_mesh_rank(device_mesh, dim):
    if device_mesh is None:
        device_mesh = get_device_mesh()

    if isinstance(device_mesh, TorchMockDeviceMesh):
        global_rank = torch.distributed.distributed_c10d.get_rank()
        device_mesh_shape = get_device_mesh().shape
        world_size = functools.reduce(operator.mul, device_mesh_shape)
        rank_coords = (numpy.arange(world_size).reshape(
            *device_mesh_shape) == global_rank).nonzero()
        rank = rank_coords[dim].item()
    elif isinstance(device_mesh, DeviceMesh):
        rank = device_mesh.get_coordinate_on_dim(dim)
    else:
        raise RuntimeError("DeviceMesh not support or not initialize")

    return rank

easydist/torch/device_mesh.py
alibaba/easydist
# Copyright (c) 2023, Alibaba Group;
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================

import logging
import importlib
from functools import update_wrapper
from typing import Any

import numpy
import torch
import torch.utils._pytree as pytree
from torch.distributed._tensor import DeviceMesh

import easydist.config as mdconfig
from easydist.torch.device_mesh import (device_mesh_world_size, get_device_mesh, set_device_mesh)
from easydist.torch.experimental.compiler import _compile
from easydist.torch.experimental.init_helper import (CpuModuleInitHelper, SetParaInitHelper)
from easydist.torch.utils import get_input_signature

logger = logging.getLogger(__name__)


class CompiledFuncWrapper:

    def __init__(self,
                 func,
                 tracing_mode="fake",
                 cuda_graph=True,
                 enable_mono_graph=False,
                 compile_only=False) -> None:
        update_wrapper(self, func)
        self.original_func = func

        self.compiled_func = None
        self.tracing_mode = tracing_mode
        self.enable_cuda_graph = cuda_graph
        self.enable_mono_graph = enable_mono_graph
        self.compile_only = compile_only

        self.init_helper = SetParaInitHelper()

        self.all_input_signature = []
        self.graph_list = {}

        self.cuda_graph_space = {}
        self.graph_pool = None

    def register_input_signature(self, *args: Any, **kwargs: Any) -> str:

        input_signature = get_input_signature(*args, **kwargs)

        if input_signature not in self.all_input_signature:

            self.all_input_signature.append(input_signature)
            logger.info(f"[Compile API] register input signature: {input_signature}")

            if self.enable_cuda_graph:
                self.cuda_graph_space[input_signature] = {
                    "cuda_graph": None,
                    "cuda_graph_input": None,
                    "cuda_graph_output": None,
                }

        return input_signature

    def register_cpu_module(self, cpu_module):
        self.init_helper = CpuModuleInitHelper(cpu_module)

    def set_init_helper(self, init_helper):
        self.init_helper = init_helper

    def __call__(self, *args: Any, **kwargs: Any) -> Any:

        input_signature = self.register_input_signature(*args, **kwargs)

        # skip compile when only one device
        world_size = torch.distributed.get_world_size()
        need_compile = world_size >= 2
        if get_device_mesh() is not None:
            need_compile = device_mesh_world_size() >= 2
        if need_compile and self.compiled_func is None:
            mesh_shape = numpy.array(range(world_size)).reshape(1, -1)
            mesh = DeviceMesh("cuda", mesh_shape.tolist())

            if get_device_mesh() == None:
                set_device_mesh(mesh)

            self.compiled_func = _compile(self.original_func, self.tracing_mode, self.init_helper,
                                          input_signature, args, kwargs)
            self.graph_list[input_signature] = self.compiled_func.graph
            # release the cpu module when finised pre-shard in _compiler
            self.init_helper = None
        if self.compile_only:
            return self.compiled_func

        def run_func(*args, **kwargs):
            if self.compiled_func:
                if input_signature not in self.graph_list:
                    if self.enable_mono_graph:
                        mono_graph = self.compiled_func.compile_mono_graph(*args, **kwargs)
                        self.graph_list[input_signature] = mono_graph
                        logger.info(f"[Compile API] compile mono graph for {input_signature}")
                    else:
                        raise RuntimeError(
                            "Input mismatch. If you are sure that different inputs do not change the graph, "
                            "you can try turning on the enable_mono_graph option.")
                graph = self.graph_list[input_signature]
                output = self.compiled_func.run_with_graph(graph, *args, **kwargs)
                if importlib.util.find_spec("torch.distributed._functional_collectives"):
                    from torch.distributed._functional_collectives import AsyncCollectiveTensor
                    # wait and unwrap AsyncCollectiveTensor
                    def wait_unwarp_fn(async_coll_tensor_):
                        async_coll_tensor_.trigger_wait()
                        return async_coll_tensor_.elem
                    return pytree.tree_map_only(AsyncCollectiveTensor, wait_unwarp_fn,output)
                else:
                    return output
            else:
                return self.original_func(*args, **kwargs)

        if self.enable_cuda_graph:
            current_space = self.cuda_graph_space[input_signature]
            if current_space["cuda_graph"] is None:
                logger.info("[Compile API] cuda graph warming up...")

                flatten_args, args_specs = pytree.tree_flatten([args, kwargs])
                current_space["cuda_graph_input"] = []
                for f_args in flatten_args:
                    if isinstance(f_args, torch.Tensor):
                        current_space["cuda_graph_input"].append(
                            torch.empty_like(f_args).copy_(f_args))
                    else:
                        current_space["cuda_graph_input"].append(f_args)
                args, kwargs = pytree.tree_unflatten(current_space["cuda_graph_input"], args_specs)

                # warm up
                s = torch.cuda.Stream()
                s.wait_stream(torch.cuda.current_stream())
                with torch.cuda.stream(s):
                    current_space["cuda_graph_output"] = run_func(*args, **kwargs)
                torch.cuda.current_stream().wait_stream(s)
                torch.cuda.synchronize()

                current_space["cuda_graph"] = torch.cuda.CUDAGraph()

                with torch.cuda.graph(current_space["cuda_graph"], self.graph_pool):
                    current_space["cuda_graph_output"] = run_func(*args, **kwargs)

                # assign the pool of first cuda graph as self.graph_pool
                if self.graph_pool is None:
                    self.graph_pool = current_space["cuda_graph"].pool()
            else:
                flatten_args, args_specs = pytree.tree_flatten([args, kwargs])
                for static_input, outside_input in zip(current_space["cuda_graph_input"],
                                                       flatten_args):
                    if isinstance(static_input, torch.Tensor):
                        static_input.copy_(outside_input)

            current_space["cuda_graph"].replay()
            return current_space["cuda_graph_output"]
        else:
            return run_func(*args, **kwargs)


def easydist_compile(func=None,
                     tracing_mode="fake",
                     cuda_graph=True,
                     enable_mono_graph=False,
                     use_hint=False,
                     liveness_only_input=False,
                     max_solver_time=float("inf"),
                     compile_only=False):

    mdconfig.use_hint = use_hint
    mdconfig.liveness_only_input = liveness_only_input
    mdconfig.max_seconds_same_incumbent = max_solver_time

    if func:
        return CompiledFuncWrapper(func, tracing_mode, cuda_graph, enable_mono_graph, compile_only)
    else:

        def wrapper(func):
            return CompiledFuncWrapper(func, tracing_mode, cuda_graph, enable_mono_graph, compile_only)

        return wrapper

easydist/torch/experimental/api.py
alibaba/easydist
# Copyright (c) 2023, Alibaba Group;
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================

import logging
import os
import pickle
import time
import threading
from functools import partial
from typing import Any, cast
from contextlib import nullcontext

import numpy
import rich
import torch
import torch.utils._pytree as pytree
import torch.distributed.rpc as rpc
from torch._subclasses.fake_tensor import FakeTensor, FakeTensorMode
from torch.distributed._tensor import (DeviceMesh, DTensor, Replicate, distribute_tensor)
from torch.fx._pytree import tree_flatten_spec
from torch.fx.experimental.proxy_tensor import make_fx
from torch.nn.utils import stateless

import easydist.config as mdconfig
from easydist.autoflow.solver import AutoFlowSolver
from easydist.metashard import metair
from easydist.torch.bridge import (get_torch_sharding_strategy, to_torch_spmd, torch2meta_graph)
from easydist.torch.experimental.decomp_utils import EASYDIST_DECOMP_TABLE
from easydist.torch.experimental.init_helper import (SetParaInitHelper, init_contiguous_buf,
                                                     materialize_zero)
from easydist.torch.passes import (eliminate_detach, fix_addmm_bias, fix_convoluation_bias,
                                   fix_embedding, fix_meta_device, sharding_transform)
from easydist.torch.device_mesh import get_device_mesh, set_device_mesh
from easydist.torch.mem_anaylize import mem_anaylize
from easydist.torch.passes import comm_optimize, rule_override_by_graph
from easydist.torch.sharding_interpreter import EDTorchShardingAnn
from easydist.torch.utils import (_enable_compile, _rematerialize_optimizer, _sharding_ann_env)
from easydist.utils import rgetattr, rsetattr
from easydist.utils.testing import TorchMockDeviceMesh

# for pickle dump opt_strategy
import sys

sys.setrecursionlimit(100000)

logger = logging.getLogger(__name__)

sharding_sol = None
sol_rdy_cond = threading.Condition()


def preprocess_traced_graph(fx_module: torch.fx.GraphModule):
    fx_module = fix_meta_device(fx_module)
    fx_module = fix_embedding(fx_module)
    fx_module = fix_addmm_bias(fx_module)
    # fx_module = fix_convoluation_bias(fx_module)
    fx_module = eliminate_detach(fx_module)

    fx_module.recompile()

    return fx_module


def easydist_shard(fx_module: torch.fx.GraphModule, state_tensor_num, input_signature, *args, **kwargs):
    # only called by rank 0
    if mdconfig.enable_compile_cache:
        os.makedirs(mdconfig.compile_cache_dir, exist_ok=True)
        compiled_cache_file = os.path.join(mdconfig.compile_cache_dir, f"./{input_signature}.pkl")

    find_cache = mdconfig.enable_compile_cache and os.path.exists(compiled_cache_file)

    if find_cache:
        logger.info(f"load compiled result from {compiled_cache_file}.")
        shape_info, opt_strategy, sharding_strategy, args_strategy = pickle.load(
            open(compiled_cache_file, "rb"))
    else:
        if mdconfig.log_level <= logging.DEBUG:
            fx_module.print_readable()

        # (1) sharding annotation
        with _sharding_ann_env():
            start_t = time.perf_counter()
            sharding_interpreter = EDTorchShardingAnn(fx_module)
            flatten_args = tree_flatten_spec(list(args) + list(kwargs.values()), fx_module._in_spec)
            sharding_info, shape_info = sharding_interpreter.run(*flatten_args)
            logger.info(f"[EDTorchShardingAnn.time]:\t {time.perf_counter() - start_t} s.")
            if mdconfig.log_level <= logging.DEBUG:
                rich.print("sharding_info:\n", sharding_info)
                rich.print("shape_info:\n", shape_info)

        # (2) translate fx.GraphModule into MetaGraph
        meta_graph = torch2meta_graph(fx_module, state_tensor_num, sharding_info, shape_info)

        if mdconfig.log_level <= logging.DEBUG:
            rich.print(meta_graph)

        # (3) construct AutoFlowSolver and run ILP
        device_mesh = get_device_mesh()
        device_mesh_shape = (device_mesh.size(0), device_mesh.size(1))

        total_memery = torch.cuda.get_device_properties(torch.cuda.current_device()).total_memory
        solver = AutoFlowSolver(device_mesh_shape, total_memery=total_memery)

        if mdconfig.enable_graph_coarsen:
            logger.info(f"enable graph coarsen with level {mdconfig.coarsen_level}.")
            solver.add_coarsen_graph(meta_graph)
        else:
            solver.add_graph(meta_graph)

        start_t = time.perf_counter()
        if mdconfig.enable_graph_coarsen:
            opt_strategy = solver.ilp_solve()
        else:
            opt_strategy = solver.ilp_optimize()
        logger.info(f"[AutoFlowSolver.time]:\t {time.perf_counter() - start_t} s.")

        if mdconfig.log_level <= logging.DEBUG:
            rich.print(opt_strategy)

        sharding_strategy = get_torch_sharding_strategy(fx_module, opt_strategy)

        if mdconfig.log_level <= logging.DEBUG:
            rich.print(sharding_strategy)

        args_strategy = meta_graph.get_input_strategy(opt_strategy)
        args_strategy = [[to_torch_spmd(i) for i in var_strategy]
                         for var_strategy in args_strategy]

        if mdconfig.enable_compile_cache:
            pickle.dump([shape_info, opt_strategy, sharding_strategy, args_strategy],
                        open(compiled_cache_file, "wb"))
            logger.info(f"compiled result saved in {compiled_cache_file}.")

    return shape_info, opt_strategy, sharding_strategy, args_strategy


@torch.no_grad()
def sharded_tensor(tensor, strategy, mesh, materialize_fn):
    # if tensor is DTensor, redistribute it
    if isinstance(tensor, DTensor):
        return tensor.redistribute(mesh, strategy)

    # materialize FakeTensor and distribute_tensor
    if isinstance(tensor, torch._subclasses.fake_tensor.FakeTensor):
        tensor = materialize_fn(tensor)
    if tensor.is_meta:
        tensor = materialize_fn(tensor)
    return distribute_tensor(tensor, mesh, strategy)


@torch.no_grad()
def dtensor_to_tensor(leaf):
    if isinstance(leaf, DTensor):
        replicate_leaf = leaf.redistribute(leaf.device_mesh, [Replicate()] * len(leaf.placements))
        return replicate_leaf.to_local()
    return leaf


def fetch_strategy():
    with sol_rdy_cond:
        if sharding_sol is None:
            sol_rdy_cond.wait()

    return sharding_sol


def _compile(func, tracing_mode, init_helper, input_signature, args, kwargs):
    module, opt = None, None

    for arg in pytree.tree_flatten(list(args) + list(kwargs.values()))[0]:
        if isinstance(arg, torch.nn.Module):
            assert module is None, "Only support single nn.Module in args now"
            module = arg
        if isinstance(arg, torch.optim.Optimizer):
            assert opt is None, "Only support single Optimizer in args now"
            opt = arg

    params, buffers = {}, {}
    if module is not None:
        params = dict(module.named_parameters())
        buffers = dict(module.named_buffers())

        if isinstance(init_helper, SetParaInitHelper):
            init_helper.module = module

    named_states = {}

    if opt is not None:
        # assign grad and warm up optimizer
        mode = nullcontext()
        for name in dict(module.named_parameters()):
            with torch.no_grad():
                rsetattr(module, name + ".grad", torch.zeros_like(rgetattr(module, name).data))
                if isinstance(rgetattr(module, name).data, FakeTensor):
                    mode = FakeTensorMode()
        with mode:
            opt.step()
            opt.zero_grad(True)

        for n, p in params.items():
            if p in opt.state:
                named_states[n] = opt.state[p]  # type: ignore[index]
                # if step in state, reduce one for warmup step.
                if 'step' in named_states[n]:
                    named_states[n]['step'] -= 1

    flat_named_states, named_states_spec = pytree.tree_flatten(named_states)

    # fix for sgd withtout momentum
    if all(state is None for state in flat_named_states):
        named_states = {}
        flat_named_states, named_states_spec = pytree.tree_flatten(named_states)

    state_tensor_num = len(params) + len(buffers) + len(flat_named_states)

    def stateless_func(func, params, buffers, named_states, args, kwargs):
        with stateless._reparametrize_module(
                cast(torch.nn.Module, module), {
                    **params,
                    **buffers
                }, tie_weights=True) if module else nullcontext(), _rematerialize_optimizer(
                    opt, named_states, params) if opt else nullcontext():
            ret = func(*args, **kwargs)

        grads = {k: v.grad for k, v in params.items()}
        return params, buffers, named_states, grads, ret

    with _enable_compile():
        traced_graph = make_fx(partial(stateless_func, func),
                               tracing_mode=tracing_mode,
                               decomposition_table=EASYDIST_DECOMP_TABLE,
                               _allow_non_fake_inputs=False)(params, buffers, named_states, args,
                                                             kwargs)

    traced_graph.graph.eliminate_dead_code()
    traced_graph = preprocess_traced_graph(traced_graph)
    traced_graph.recompile()

    world_size = torch.distributed.get_world_size()
    rank = torch.distributed.get_rank()

    # Lansong(TODO) Currently send strategy by rpc. But broadcast way is more efficient.
    rpc.init_rpc(f"ed_worker{rank}", rank=rank, world_size=world_size)
    if rank == 0:
        shape_info, opt_strategy, sharding_strategy, args_strategy = easydist_shard(
                                traced_graph, state_tensor_num, input_signature,
                                params, buffers, named_states, args, kwargs)

        with sol_rdy_cond:
            global sharding_sol
            sharding_sol = [shape_info, opt_strategy, sharding_strategy, args_strategy]
            sol_rdy_cond.notify_all()
    else:
        shape_info, opt_strategy, sharding_strategy, args_strategy = rpc.rpc_sync(
                              "ed_worker0", fetch_strategy, args=(), timeout=0)

    rpc.shutdown()

    # mem_anaylize(traced_graph, shape_info, opt_strategy)
    sharded_graph = sharding_transform(traced_graph, sharding_strategy)
    sharded_graph = fix_embedding(sharded_graph, recover=True)

    if mdconfig.comm_optimization is True:
        sharded_graph = comm_optimize(sharded_graph, opt_strategy)

    # override pytorch dtensor propagate rules to optimize dispater behavior
    if mdconfig.override_dtensor_rule is True:
        sharded_graph = rule_override_by_graph(sharded_graph, opt_strategy, shape_info)

    if mdconfig.log_level <= logging.DEBUG:
        sharded_graph.print_readable()

    # do not use mock device after get sharded_graph
    device_mesh = get_device_mesh()
    if isinstance(device_mesh, TorchMockDeviceMesh):
        if device_mesh.debug_only:
            mesh_shape = numpy.array(range(world_size)).reshape(1, -1).tolist()
        else:
            mesh_shape = device_mesh.shape
        mesh = DeviceMesh("cuda", mesh_shape)
        set_device_mesh(mesh)

    device = mdconfig.easydist_device

    # keep fake params, buffers, named_states
    fake_tensor_mode = FakeTensorMode()

    def wrap_fake(x):
        if isinstance(x, torch.Tensor):
            return fake_tensor_mode.from_tensor(x)
        return x

    fake_params = pytree.tree_map(wrap_fake, params)
    fake_buffers = pytree.tree_map(wrap_fake, buffers)
    fake_named_states = pytree.tree_map(wrap_fake, named_states)

    # pre-shard params, buffers, named_states
    params_strategy = args_strategy[:len(params)]
    buffers_strategy = args_strategy[len(params):len(params) + len(buffers)]

    if mdconfig.use_contiguous_buffer:
        contiguous_buf = init_contiguous_buf(params, params_strategy, device_mesh)

    index = 0
    for idx, param_name in enumerate(params):
        materialize_fn = init_helper.get_materialize_fn()
        materialize_fn = partial(materialize_fn,
                                 param_buf_key=param_name,
                                 materialization_device=device)
        params[param_name] = sharded_tensor(params[param_name],
                                            params_strategy[idx],
                                            get_device_mesh(),
                                            materialize_fn=materialize_fn)

        size = params[param_name]._local_tensor.numel()

        if mdconfig.use_contiguous_buffer:
            contiguous_buf[index:index + size] = params[param_name]._local_tensor.view(-1)
            params[param_name]._local_tensor = contiguous_buf[index:index + size].view(
                params[param_name]._local_tensor.shape)

        index += size

    for idx, buffer_name in enumerate(buffers):
        materialize_fn = init_helper.get_materialize_fn()
        materialize_fn = partial(materialize_fn,
                                 param_buf_key=buffer_name,
                                 materialization_device=device)
        buffers[buffer_name] = sharded_tensor(buffers[buffer_name],
                                              buffers_strategy[idx],
                                              get_device_mesh(),
                                              materialize_fn=materialize_fn)

    # use zero init for optimizer states
    flat_named_states, named_states_spec = pytree.tree_flatten(named_states)
    state_tensor_num = len(params) + len(buffers)
    materialize_fn = partial(materialize_zero, materialization_device=device)
    for i in range(len(flat_named_states)):
        if isinstance(flat_named_states[i], torch.Tensor):
            flat_named_states[i] = sharded_tensor(flat_named_states[i],
                                                  args_strategy[state_tensor_num],
                                                  get_device_mesh(),
                                                  materialize_fn=materialize_fn)
            state_tensor_num += 1

    named_states = pytree.tree_unflatten(flat_named_states, named_states_spec)


    class EDCompiledFunc:

        def __init__(self, graph) -> None:
            self.graph = graph

        @torch.no_grad()
        def compiled_func(self, graph, *args, **kwargs):
            nonlocal params, buffers, named_states

            # tensor to DTensor
            # (TODO) the redistributed of params, buffers, named_states should in graph
            flatten_args, args_specs = pytree.tree_flatten([args, kwargs])

            device = mdconfig.easydist_device
            materialize_fn = partial(materialize_zero, materialization_device=device)

            args_strategy_idx = state_tensor_num
            for i in range(len(flatten_args)):
                if isinstance(flatten_args[i], torch.Tensor):
                    flatten_args[i] = sharded_tensor(flatten_args[i].detach(),
                                                     args_strategy[args_strategy_idx],
                                                     get_device_mesh(),
                                                     materialize_fn=materialize_fn)
                    args_strategy_idx += 1

            args, kwargs = pytree.tree_unflatten(flatten_args, args_specs)

            # run the sharded_graph
            params, buffers, named_states, grads, sharded_out = graph(params, buffers, named_states,
                                                                      args, kwargs)

            for para_name in params:
                params[para_name].grad = grads[para_name]

            # out from DTensor to Tensor
            local_out = pytree.tree_map(dtensor_to_tensor, sharded_out)

            return local_out

        def __call__(self, *args: Any, **kwargs: Any) -> Any:
            return self.compiled_func(self.graph, *args, **kwargs)

        def run_with_graph(self, graph, *args: Any, **kwargs: Any) -> Any:
            return self.compiled_func(graph, *args, **kwargs)

        def get_state(self):
            return params, buffers, named_states

        def compile_mono_graph(self, *args: Any, **kwargs: Any):
            with _enable_compile():
                fx_module = make_fx(partial(stateless_func, func),
                                    tracing_mode=tracing_mode,
                                    decomposition_table=EASYDIST_DECOMP_TABLE,
                                    _allow_non_fake_inputs=False)(fake_params, fake_buffers,
                                                                  fake_named_states, args, kwargs)

            fx_module.graph.eliminate_dead_code()
            fx_module.recompile()

            fx_module = preprocess_traced_graph(fx_module)

            sharded_fx_module = sharding_transform(fx_module, sharding_strategy)
            sharded_fx_module = fix_embedding(sharded_fx_module, recover=True)

            if mdconfig.log_level <= logging.DEBUG:
                sharded_fx_module.print_readable()

            return sharded_fx_module

        def parameters(self):
            return params.values()

        def named_parameters(self):
            return params

    # release all cuda memory from module here
    # the param maintain in the local of compiled function.
    if module is not None:
        module.to("meta")

    return EDCompiledFunc(sharded_graph)

easydist/torch/experimental/compiler.py
alibaba/easydist
# Most functions in this file are copied from https://github.com/pytorch/pytorch/blob/main/torch/distributed/_spmd/api.py.
# We copy these functions because they are not ready in released pytorch.

from functools import partial

import torch
from torch._decomp.decompositions import mse_loss, mse_loss_backward

aten = torch.ops.aten  # pyre-ignore


# from torch/distributed/_spmd/api.py
def _fused_adam_decomp(
    self,
    grads,
    exp_avgs,
    exp_avg_sqs,
    max_exp_avg_sqs,
    state_steps,
    *,
    lr=1,
    beta1=1,
    beta2=1,
    weight_decay=1,
    eps=1,
    amsgrad=True,
    maximize=True,
    grad_scale=None,
    found_inf=None,
):
    orig_tuple = (self, grads, exp_avgs, exp_avg_sqs, max_exp_avg_sqs)
    updated_tuple = aten._fused_adam.default(
        self,
        grads,
        exp_avgs,
        exp_avg_sqs,
        max_exp_avg_sqs,
        state_steps,
        lr=lr,
        beta1=beta1,
        beta2=beta2,
        weight_decay=weight_decay,
        eps=eps,
        amsgrad=amsgrad,
        maximize=maximize,
        grad_scale=grad_scale,
        found_inf=found_inf,
    )

    for idx, (orig, updated) in enumerate(zip(orig_tuple, updated_tuple)):
        if idx == 1:
            # skip gradient copying as we don't need to copy gradients back
            continue
        for o, u in zip(orig, updated):
            o.copy_(u)


# from torch/distributed/_spmd/api.py
def _foreach_add_decomp(self, other, alpha=1):
    self_updated = aten._foreach_add.List(self, other, alpha=alpha)
    for s, s_u in zip(self, self_updated):
        s.copy_(s_u)


# from torch/distributed/_spmd/api.py
def _foreach_unaop_decomp(op, self):
    self_updated = op(self)
    for s, s_u in zip(self, self_updated):
        s.copy_(s_u)


# from torch/distributed/_spmd/api.py
def _foreach_binop_list_decomp(op, self, other):
    self_updated = op(self, other)
    for s, s_u in zip(self, self_updated):
        s.copy_(s_u)


# from torch/distributed/_spmd/api.py
def _foreach_triop_scalar_decomp(op, self, other, scalar):
    self_updated = op(self, other, scalar)
    for s, s_u in zip(self, self_updated):
        s.copy_(s_u)


# from torch/distributed/_spmd/api.py
def _foreach_binop_scalar_decomp(op, self, scalar=1):
    self_updated = op(self, scalar)
    for s, s_u in zip(self, self_updated):
        s.copy_(s_u)


# from torch/distributed/_spmd/api.py
def _foreach_addcop_scalar_decomp(op, self, tensor1, tensor2, scalar=1):
    self_updated = op(self, tensor1, tensor2, scalar)
    for s, s_u in zip(self, self_updated):
        s.copy_(s_u)


# from torch/distributed/_spmd/api.py
def _foreach_addcop_tensor_decomp(op, self, tensor1, tensor2, tensor):
    self_updated = op(self, tensor1, tensor2, tensor)
    for s, s_u in zip(self, self_updated):
        s.copy_(s_u)


# modified from torch/distributed/_spmd/api.py
EASYDIST_DECOMP_TABLE = {
    aten._foreach_add_.List: _foreach_add_decomp,
    aten._foreach_add_.Scalar: partial(_foreach_binop_scalar_decomp, aten._foreach_add.Scalar),
    aten._foreach_addcdiv_.Scalar: partial(_foreach_addcop_scalar_decomp,
                                           aten._foreach_addcdiv.Scalar),
    aten._foreach_addcdiv_.Tensor: partial(_foreach_addcop_tensor_decomp,
                                           aten._foreach_addcdiv.Tensor),
    aten._foreach_addcmul_.Scalar: partial(_foreach_addcop_scalar_decomp,
                                           aten._foreach_addcmul.Scalar),
    aten._foreach_div_.List: partial(_foreach_binop_list_decomp, aten._foreach_div.List),
    aten._foreach_mul_.Scalar: partial(_foreach_binop_list_decomp, aten._foreach_mul.Scalar),
    aten._foreach_div_.Scalar: partial(_foreach_binop_scalar_decomp, aten._foreach_div.Scalar),
    aten._foreach_neg_.default: partial(_foreach_unaop_decomp, aten._foreach_neg.default),
    aten._foreach_lerp_.Scalar: partial(_foreach_triop_scalar_decomp, aten._foreach_lerp.Scalar),
    aten._foreach_reciprocal_.default: partial(_foreach_unaop_decomp,
                                               aten._foreach_reciprocal.default),
    aten._foreach_sqrt_.default: partial(_foreach_unaop_decomp, aten._foreach_sqrt.default),
    aten._foreach_sub_.Scalar: partial(_foreach_binop_scalar_decomp, aten._foreach_sub.Scalar),
    aten._fused_adam_.default: _fused_adam_decomp,
    aten.mse_loss.default: mse_loss,
    aten.mse_loss_backward.default: mse_loss_backward,
}

try:
    from torch._inductor.decomposition import _foreach_lerp_scalar
    EASYDIST_DECOMP_TABLE.update({aten._foreach_lerp.Scalar: _foreach_lerp_scalar})
except:
    pass

easydist/torch/experimental/decomp_utils.py
alibaba/easydist
# Copyright (c) 2023, Alibaba Group;
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================

from functools import partial

import torch
import torch.utils._pytree as pytree
from torch._subclasses import FakeTensor
from torch.utils._mode_utils import no_dispatch


def fake_to_real(e):
    with no_dispatch():
        if isinstance(e, FakeTensor):
            out = torch.zeros_like(e, device=e.fake_device)
            return out
        return e


@torch.no_grad()
def materialize_zero(tensor, materialization_device):
    tensor = fake_to_real(tensor)
    return torch.zeros_like(tensor, device=materialization_device)


@torch.no_grad()
def materialize_random(tensor, materialization_device):
    if tensor.dtype == torch.bool:
        return torch.rand(tensor.size(), dtype=torch.float, device=materialization_device) > 0.5
    elif torch.is_floating_point(tensor):
        return torch.rand(tensor.size(), dtype=tensor.dtype, device=materialization_device)
    else:
        return torch.randint(high=8,
                             size=tensor.size(),
                             dtype=tensor.dtype,
                             device=materialization_device)


@torch.no_grad()
def materialize_module_from_cpu(tensor, param_buf_key, cpu_model: torch.nn.Module,
                                materialization_device):
    cpu_params_buffers = {
        **dict(cpu_model.named_parameters()),
        **dict(cpu_model.named_buffers()),
    }
    if param_buf_key in cpu_params_buffers:
        return cpu_params_buffers[param_buf_key].to(device=materialization_device)

    return materialize_zero(tensor, materialization_device)


@torch.no_grad()
def materialize_module_from_set_parameter(tensor, param_buf_key, module: torch.nn.Module,
                                          materialization_device):
    tensor = fake_to_real(tensor)

    submodule_name_list = param_buf_key.split(".")

    if "reset_parameters" in module.__dir__():
        if param_buf_key in dict(module.named_parameters()):
            ori_parameters = module._parameters
            module._parameters = pytree.tree_map(fake_to_real, module._parameters)
            module.to_empty(device=materialization_device)
            module.reset_parameters()
            para = dict(module.named_parameters()).get(param_buf_key, None)
            module._parameters = ori_parameters
            return para

    if "reset_buffers" in module.__dir__():
        if param_buf_key in dict(module.named_buffers()):
            ori_buffers = module._buffers
            module._buffers = pytree.tree_map(fake_to_real, module._buffers)
            module.reset_buffers()
            buf = dict(module.named_buffers()).get(param_buf_key, None)
            module._buffers = ori_buffers
            return buf

    if len(submodule_name_list) == 1:
        return materialize_zero(tensor, materialization_device)

    submodule_name, sub_param_buf_key = submodule_name_list[0], ".".join(submodule_name_list[1:])

    if submodule_name in module._modules:
        return materialize_module_from_set_parameter(tensor, sub_param_buf_key,
                                                     module._modules[submodule_name],
                                                     materialization_device)

    return materialize_zero(tensor, materialization_device)


class InitHelper:

    def __init__(self) -> None:
        pass

    def get_materialize_fn(self):
        pass


class SetParaInitHelper:

    def __init__(self, module=None) -> None:
        self.module = module

    def get_materialize_fn(self):
        return partial(materialize_module_from_set_parameter, module=self.module)


class CpuModuleInitHelper:

    def __init__(self, cpu_module) -> None:
        self.cpu_module = cpu_module

    def get_materialize_fn(self):
        return partial(materialize_module_from_cpu, cpu_model=self.cpu_module)


def init_contiguous_buf(params, params_strategy, device_mesh):

    params = list(params.values())
    dtype = params[0].dtype
    device = params[0].device
    if not all(p.dtype == dtype for p in params):
        raise ValueError("All parameters must be of the same dtype.")
    if not all(p.device == device for p in params):
        raise ValueError("All parameters must be on the same device.")
    contiguous_buf_size = 0
    assert len(params_strategy) == len(params), "mismatch length of params_strategy and params"
    for p_s, p in zip(params_strategy, params):
        tensor_size = p.numel()
        for mesh_dim_idx, placement in enumerate(p_s):
            if placement.is_shard():
                tensor_size = (tensor_size // device_mesh.mesh.shape[mesh_dim_idx]) + 1

        contiguous_buf_size += tensor_size

    return torch.empty(contiguous_buf_size, dtype=dtype, device=device)

easydist/torch/experimental/init_helper.py
alibaba/easydist
# Copyright (c) 2023, Alibaba Group;
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================

from functools import reduce
from operator import mul
from typing import Dict, List

import torch
from torch.fx.node import Node, _get_qualified_name, map_arg

from .bridge import to_torch_spmd
from .device_mesh import device_mesh_world_size

DTYPE_BYTE = {
    torch.float32: 4,
    torch.float16: 2,
    torch.int32: 4,
    torch.int64: 8,
    torch.uint8: 1,
    torch.bool: 0.125
}


def tensor_mem_size(shape_info, mem_ratio):
    mem_size = 0
    if isinstance(shape_info, tuple):
        for shape_info_, mem_ratio_ in zip(shape_info, mem_ratio):
            if 'shape' in shape_info_:
                if len(shape_info_['shape']) == 0:
                    mem_size += DTYPE_BYTE[shape_info_['dtype']]
                else:
                    mem_size += reduce(
                        mul, shape_info_['shape']) * DTYPE_BYTE[shape_info_['dtype']] * mem_ratio_
    else:
        if 'shape' in shape_info:
            if len(shape_info['shape']) == 0:
                mem_size += DTYPE_BYTE[shape_info['dtype']]
            else:
                mem_size = reduce(
                    mul, shape_info['shape']) * DTYPE_BYTE[shape_info['dtype']] * mem_ratio
    return mem_size


def _get_input_strategy(fx_graph, opt_strategy):
    input_strategy = {}

    input_name = []
    for node in fx_graph.graph.nodes:
        if node.op == 'placeholder':
            input_name.append(node.name)

    for node in reversed(fx_graph.graph.nodes):
        if node.op == "call_function":
            op_name = _get_qualified_name(node.target)
            if op_name != "_operator.getitem":
                unique_key = node.name
                if unique_key in opt_strategy:
                    invars = opt_strategy[unique_key]['node'].invars
                    for idx, var in enumerate(invars):
                        if var.name in input_name:
                            input_strategy[var.name] = [
                                to_torch_spmd(i)
                                for i in opt_strategy[unique_key]['strategy'].in_strtg_group[idx]
                            ]

    return input_strategy


def mem_anaylize(fx_graph: torch.fx.GraphModule, tensor_shape_info, opt_strategy):

    same_mem_funcs = [
        "torch.ops.aten.t", "torch.ops.aten.view", "torch.ops.aten._unsafe_view",
        "torch.ops.aten.slice", "torch.ops.aten.transpose", "torch.ops.aten.squeeze",
        "torch.ops.aten.unsqueeze", "torch.ops.aten.permute", "torch.ops.aten.slice",
        "_operator.getitem", "torch.ops.aten.expand"
    ]

    # Run through reverse nodes and record the first instance of a use
    # of a given node. This represents the *last* use of the node in the
    # execution order of the program, which we will use to free unused
    # values
    node_to_last_use: Dict[Node, Node] = {}
    user_to_last_uses: Dict[Node, List[Node]] = {}

    def register_last_uses(n: Node, user: Node):
        if n not in node_to_last_use:
            node_to_last_use[n] = user
            user_to_last_uses.setdefault(user, []).append(n)

    for node in reversed(fx_graph.graph.nodes):
        map_arg(node.args, lambda n: register_last_uses(n, node))
        map_arg(node.kwargs, lambda n: register_last_uses(n, node))

    persistent_mem, activation_mem = {}, {}
    mem_ratio = {}
    same_memory_tensor_info = {}

    total_memory, peak_memory = 0, 0

    f = open("mem_info.txt", "w")

    world_size = device_mesh_world_size()

    input_strategy = _get_input_strategy(fx_graph, opt_strategy)

    for node in fx_graph.graph.nodes:
        # input and parametes will not release during inference
        if node.op == 'placeholder':

            mem_ratio[node.name] = 1
            if node.name in input_strategy:
                for place in input_strategy[node.name]:
                    if place.is_shard():
                        mem_ratio[node.name] = 1 / world_size

            persistent_mem[node.name] = tensor_mem_size(tensor_shape_info[node.name],
                                                        mem_ratio[node.name])
            total_memory += persistent_mem[node.name]
            if peak_memory < total_memory:
                peak_memory = total_memory

        if node.op == 'call_function':
            if _get_qualified_name(node.target) in same_mem_funcs:
                if node.args[0].name not in persistent_mem:
                    if node.args[0] not in same_memory_tensor_info:
                        same_memory_tensor_info[node.args[0].name] = [node.args[0].name]
                    same_memory_tensor_info[node.args[0].name].append(node.name)
            else:
                mem_ratio[node.name] = [1]
                if node.name in opt_strategy:
                    out_strategy = [[to_torch_spmd(i) for i in ii]
                                    for ii in opt_strategy[node.name]['strategy'].out_strtg_group]
                    mem_ratio[node.name] = [1] * len(out_strategy)
                    for idx, out_stgy in enumerate(out_strategy):
                        for place in out_stgy:
                            if place.is_shard():
                                mem_ratio[node.name][idx] = 1 / world_size

                if not isinstance(tensor_shape_info[node.name], tuple):
                    mem_ratio[node.name] = mem_ratio[node.name][0]

                activation_mem[node.name] = tensor_mem_size(tensor_shape_info[node.name],
                                                            mem_ratio[node.name])
                total_memory += activation_mem[node.name]
                if peak_memory < total_memory:
                    peak_memory = total_memory

        # release used tensor
        for to_delete in user_to_last_uses.get(node, []):

            for value in same_memory_tensor_info.values():
                while to_delete.name in value:
                    value.remove(to_delete.name)

            new_same_memory_tensor_info = {}
            for k, v in same_memory_tensor_info.items():
                if len(v) > 0:
                    new_same_memory_tensor_info[k] = v
                else:
                    if k in activation_mem:
                        total_memory -= tensor_mem_size(tensor_shape_info[k], mem_ratio[k])
                        del activation_mem[k]
            same_memory_tensor_info = new_same_memory_tensor_info

            if to_delete.name in activation_mem:
                if to_delete.name not in same_memory_tensor_info:
                    total_memory -= tensor_mem_size(tensor_shape_info[to_delete.name],
                                                    mem_ratio[to_delete.name])
                    del activation_mem[to_delete.name]

        gb_ratio = 1024 * 1024 * 1024
        f.write(
            f"Node | Total | Peak Memory: {node} | {total_memory / gb_ratio:.4f} GB | {peak_memory / gb_ratio:.4f} GB\n"
        )

    f.close()

easydist/torch/mem_anaylize.py
alibaba/easydist
# Copyright (c) 2023, Alibaba Group;
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================

from .fix_embedding import fix_embedding
from .fix_bias import fix_addmm_bias, fix_convoluation_bias
from .fix_meta_device import fix_meta_device
from .eliminate_detach import eliminate_detach
from .sharding import sharding_transform
from .comm_optimize import comm_optimize
from .rule_override import rule_override_by_graph

__all__ = [
    "fix_embedding", "fix_addmm_bias", "fix_convoluation_bias", "eliminate_detach",
    "sharding_transform", "fix_meta_device", "comm_optimize", "rule_override_by_graph"
]

easydist/torch/passes/__init__.py
alibaba/easydist
# Copyright (c) 2023, Alibaba Group;
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================

import logging
import operator

import torch
from sortedcontainers import SortedList
from torch.distributed._tensor.api import DTensor
from torch.fx.node import _get_qualified_name

from easydist.torch.device_mesh import get_device_mesh

logger = logging.getLogger(__name__)


def comm_optimize(fx_module: torch.fx.GraphModule, opt_strategy=None, grouping=False):
    fx_module.graph.eliminate_dead_code()
    fx_module.recompile()

    node_to_rank: dict[torch.fx.Node, int] = {}
    # processing
    rank = 0
    for node in fx_module.graph.nodes:
        if node.op == 'call_function':
            op_name = _get_qualified_name(node.target)
            if op_name == 'easydist.torch.passes.sharding.redist_tensor_func':
                continue
        node_to_rank[node] = rank
        rank += 1

    '''
        ppredecesor -> predecesor1 -> node1
            |    ...                    ^
            v                           |
    predecesorN -> nodeN           predecesor1'
    '''

    # comm_op expressed as (from_node, comm_node, to_node)

    comm_queue = SortedList(key=lambda op_tup: node_to_rank[op_tup[0]])

    for node in fx_module.graph.nodes:
        if node.op == 'call_function':
            op_name = _get_qualified_name(node.target)

            if op_name != 'easydist.torch.passes.sharding.redist_tensor_func':
                input_nodes = node.all_input_nodes
                for predecesor in input_nodes:
                    if predecesor.op == 'call_function':
                        pre_op_name = _get_qualified_name(predecesor.target)
                        if pre_op_name == 'easydist.torch.passes.sharding.redist_tensor_func':
                            ppredecesor = predecesor.all_input_nodes
                            assert (len(ppredecesor) == 1)
                            from_node = ppredecesor[0]
                            comm_node = predecesor
                            to_node = node

                            # code needed to eliminate redundant comm node
                            if opt_strategy is not None and from_node.op == 'call_function':
                                target_node = from_node
                                idx = 0
                                if from_node.name.__contains__('getitem'):
                                    pppredecesor = from_node.all_input_nodes
                                    assert(len(pppredecesor) == 1)
                                    target_node = pppredecesor[0]
                                    idx = from_node.args[1]
                                to_node_strategy = opt_strategy[to_node.name]['strategy']
                                from_node_strategy = opt_strategy[target_node.name]['strategy']
                                if from_node_strategy is not None and \
                                    from_node_strategy.get_outvar_strtg(idx) == \
                                    to_node_strategy.get_invar_strtg(input_nodes.index(comm_node)):
                                    to_node.replace_input_with(comm_node, from_node)
                                    continue
                            comm_queue.add((from_node, comm_node, to_node))

    # node just computed -> commnications followed
    comm_map = {}

    # can be significantly improved
    for (from_node, comm_node, to_node) in comm_queue:
        if comm_map.get(from_node) is None:
            comm_map[from_node] = []
        comm_map[from_node].append((from_node, comm_node, to_node))

    def redist_tensor_func_transformed(input_tensors: list, input_specs: list):
        res = []
        device_mesh = get_device_mesh()
        res_cache: dict[DTensor, list] = {}
        for input_tensor, spec in zip(input_tensors, input_specs):
            if isinstance(input_tensor, DTensor) and input_tensor.size() != torch.Size([0]):
                if spec != input_tensor._spec.placements:
                    hist = res_cache.get(input_tensor)
                    current_res = None
                    if hist != None:
                        for hist_spec, idx in hist:
                            # (CAUTION) assuming no inplace operation in graph 
                            if hist_spec == spec:
                                current_res = res[idx]
                                break
                    else:
                        res_cache[input_tensor] = []
                    if current_res is None:
                        current_res = input_tensor.redistribute(
                            device_mesh, spec).contiguous()
                        res_cache[input_tensor].append((spec, len(res)))
                    res.append(current_res)
                    continue
            res.append(input_tensor.contiguous())
        return res

    # add new comm node after nodes that need comms after computation
    for node in comm_map:
        comm_list = comm_map[node]
        # redundancy remained
        input_nodes = [n for (n, _, _) in comm_list]
        input_specs = [n.args[1] for (_, n, _) in comm_list]
        with fx_module.graph.inserting_after(node):
            new_node = fx_module.graph.call_function(
                redist_tensor_func_transformed, args=(input_nodes, input_specs))
        for idx, (_, comm_node, to_node) in enumerate(comm_list):
            with fx_module.graph.inserting_after(new_node):
                idx_node = fx_module.graph.call_function(operator.getitem, args=(new_node, idx))
            to_node.replace_input_with(comm_node, idx_node)

    # at this point all old comm operators should be eliminated
    fx_module.graph.eliminate_dead_code()
    fx_module.recompile()

    logger.info("Communication Optimization: Done!")
    return fx_module

easydist/torch/passes/comm_optimize.py
alibaba/easydist
# Copyright (c) 2023, Alibaba Group;
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================

import torch.fx as fx
from torch.fx.node import _get_qualified_name


def eliminate_detach(fx_graph: fx.GraphModule):
    for node in fx_graph.graph.nodes:
        if node.op == 'call_function':
            if _get_qualified_name(node.target) == 'torch.ops.aten.detach.default':
                node.replace_all_uses_with(node.args[0])

    fx_graph.graph.eliminate_dead_code()

    return fx_graph

easydist/torch/passes/eliminate_detach.py
alibaba/easydist
# Copyright (c) 2023, Alibaba Group;
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================

import torch
from torch.fx.node import _get_qualified_name


def fix_addmm_bias(fx_module: torch.fx.GraphModule):

    for node in fx_module.graph.nodes:
        if node.op == 'call_function':
            if "torch.ops.aten.addmm.default" in _get_qualified_name(node.target):
                node.target = torch.ops.aten.mm.default
                bias = node.args[0]
                node.args = (node.args[1], node.args[2])

                with fx_module.graph.inserting_after(node):
                    add_bias_node = fx_module.graph.call_function(torch.ops.aten.add.Tensor,
                                                                  args=(node, bias))

                    node.replace_all_uses_with(add_bias_node)

                    add_bias_node.update_arg(0, node)

    fx_module.recompile()

    return fx_module


def fix_convoluation_bias(fx_module: torch.fx.GraphModule):

    for node in fx_module.graph.nodes:
        if node.op == 'call_function':
            if "torch.ops.aten.convolution.default" in _get_qualified_name(node.target):
                if node.args[2] is not None:
                    node.target = torch.ops.aten.convolution.default
                    bias = node.args[2]
                    node.args = (node.args[0], node.args[1], None, *node.args[3:])

                    with fx_module.graph.inserting_after(node):
                        bias_new = fx_module.graph.call_function(torch.ops.aten.view.default,
                                                                 args=(bias, [1, -1, 1, 1]))

                    with fx_module.graph.inserting_after(bias_new):
                        add_bias_node = fx_module.graph.call_function(torch.ops.aten.add.Tensor,
                                                                      args=(node, bias_new))

                        node.replace_all_uses_with(add_bias_node)

                        add_bias_node.update_arg(0, node)

    fx_module.recompile()

    return fx_module

easydist/torch/passes/fix_bias.py
alibaba/easydist
# Copyright (c) 2023, Alibaba Group;
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================

import torch
from torch.fx.node import _get_qualified_name


def md_embedding(weight, indices, padding_idx=-1, scale_grad_by_freq=False, sparse=False):
    if int(torch.max(indices).item()) >= weight.shape[0]:
        raise RuntimeError("embedding indice overflow")
    return torch.ops.aten.embedding.default(weight, indices, padding_idx, scale_grad_by_freq,
                                            sparse)


def fix_embedding(fx_module: torch.fx.GraphModule, recover=False):

    for node in fx_module.graph.nodes:
        if node.op == 'call_function':
            if "torch.ops.aten.embedding.default" in _get_qualified_name(node.target):
                node.target = md_embedding

            if recover and "md_embedding" in _get_qualified_name(node.target):
                node.target = torch.ops.aten.embedding.default

    fx_module.recompile()

    return fx_module

easydist/torch/passes/fix_embedding.py
alibaba/easydist
# Copyright (c) 2023, Alibaba Group;
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================

import os
import copy

import torch

import easydist.config as mdconfig


def fix_meta_device(fx_module: torch.fx.GraphModule):

    for node in fx_module.graph.nodes:
        if node.op == 'call_function':
            if "device" in node.kwargs:
                new_kwargs = dict(copy.deepcopy(node.kwargs))
                device = mdconfig.easydist_device
                new_kwargs["device"] = torch.device(device=device)
                assert isinstance(new_kwargs, dict)
                node.kwargs = new_kwargs

    fx_module.recompile()

    return fx_module

easydist/torch/passes/fix_meta_device.py
alibaba/easydist
# Copyright (c) 2023, Alibaba Group;
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================

import logging
from typing import Dict, Tuple

import torch
import torch.distributed.distributed_c10d as c10d
from torch.distributed._tensor.api import DTensor
from torch.distributed._tensor.op_schema import OpSchema, OutputSharding
from torch.distributed._tensor.placement_types import (DTensorSpec, Replicate,
                                                       Shard, _Partial)
from torch.fx.node import _get_qualified_name

from easydist.torch.device_mesh import get_device_mesh
from easydist.torch.utils import get_dtensor_spec

# used to select proper output for the same function
# [CAUTION] only support single static graph
global graph_cnt
graph_cnt: Dict[str, Tuple[int, int]] = {}
global graph_output_struct
graph_output_struct: Dict[str, str] = {}

logger = logging.getLogger(__name__)

aten = torch.ops.aten


def _transform_to_Placemnet(varstrtg):
    res = []
    for strtg in varstrtg:
        if strtg.is_replicate():
            res.append(Replicate())
        elif strtg.is_shard():
            res.append(Shard(strtg.args['dim']))
        elif strtg.is_partial():
            res.append(_Partial(eval('c10d.' + strtg.args['ops'].__str__())))
        else:
            RuntimeError('Rule_override: strtg not recognized')
    return tuple(res)


def rule_override_by_graph(fx_module: torch.fx.GraphModule, opt_strategy, shape_info):

    def operator_rule_factory(op_name, graph_rules):

        def operator_rule(op_schema: OpSchema) -> OutputSharding:
            current_cnt, lim = graph_cnt[op_name]

            mesh = get_device_mesh()
            dtensor_specs = []
            tot_kwargs = graph_rules[op_name][current_cnt]

            for kwargs in tot_kwargs:
                if kwargs is None:
                    dtensor_specs.append(None)
                    continue
                kwargs['mesh'] = mesh
                dtensor_specs.append(get_dtensor_spec(**kwargs))

            struct = graph_output_struct[op_name]
            if struct == 'elem':
                res_specs = dtensor_specs[0]
            elif struct == 'list':
                res_specs = list(dtensor_specs)
            elif struct == 'tuple':
                res_specs = tuple(dtensor_specs)

            graph_cnt[op_name] = ((current_cnt + 1) % lim, lim)
            return OutputSharding(output_spec=res_specs)

        return operator_rule

    graph_rules: dict[str, list] = {}

    for node in fx_module.graph.nodes:
        if node.op == 'call_function':
            op_name = _get_qualified_name(node.target)

            if node.name.__contains__('redist_tensor_func') \
                    or node.name.__contains__('getitem'):
                continue
            strategy = opt_strategy[node.name]['strategy']
            output_spec_list = []

            output_shapes = shape_info[node.name]
            if not isinstance(output_shapes, tuple):
                if isinstance(output_shapes, list):
                    graph_output_struct[op_name] = 'list'
                    tot_info = tuple(output_shapes)
                else:
                    graph_output_struct[op_name] = 'elem'
                    tot_info = (output_shapes,)
            else:
                graph_output_struct[op_name] = 'tuple'
                tot_info = output_shapes

            tensor_args_idx = 0
            assert len(strategy.out_strtg_group) == len(tot_info)
            for var_strtg, tensor_info in zip(strategy.out_strtg_group,
                                              tot_info):
                if tensor_info == {} or var_strtg is None:
                    output_spec_list.append(None)
                    continue
                spec = _transform_to_Placemnet(var_strtg.var_spmd_strategy)
                shape = tensor_info['shape']
                cur_kwargs = {
                    'placements': spec,
                    'ndim': len(shape),
                    'shape': shape
                }
                output_spec_list.append(cur_kwargs)

            if graph_rules.get(op_name) is None:
                graph_rules[op_name] = []
                graph_cnt[op_name] = (0, 0)

            graph_rules[op_name].append(output_spec_list)

            graph_cnt[op_name] = (0, graph_cnt[op_name][1] + 1)

    for op_name in graph_rules:
        DTensor._propagator.register_sharding_prop_rule(
            eval(op_name),
            operator_rule_factory(op_name, graph_rules)
        )

    return fx_module

easydist/torch/passes/rule_override.py
alibaba/easydist
# Copyright (c) 2023, Alibaba Group;
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================

import logging

import torch
import torch.utils._pytree as pytree
from torch.distributed._tensor import DTensor, Replicate
from torch.fx.node import Node, _get_qualified_name, map_arg

from easydist.torch.device_mesh import device_mesh_rank, get_device_mesh
from easydist.utils.testing import MockDeviceMesh

logger = logging.getLogger(__name__)


def materialize(x, device):
    if isinstance(x, torch.Tensor):
        if x.dtype == torch.bool:
            return torch.rand(x.size(), dtype=torch.float, device=device) > 0.5
        elif torch.is_floating_point(x):
            return torch.rand(x.size(), dtype=x.dtype, device=device)
        else:
            return torch.randint(high=8, size=x.size(), dtype=x.dtype, device=device)
    return x


def redist_tensor_func(input_tensor, specs):
    if isinstance(input_tensor, DTensor) and input_tensor.size() != torch.Size([0]):
        device_mesh = get_device_mesh()
        if specs != input_tensor._spec.placements:
            return input_tensor.redistribute(device_mesh, specs).contiguous()
    return input_tensor.contiguous()


def insert_spmd_head(body):
    return ["from torch.distributed._tensor import Shard, Replicate\n", *body]


@torch.no_grad()
def to_dtensor(input_tensor, placements=None, global_size=None):
    if isinstance(input_tensor, torch.Tensor) and not isinstance(input_tensor, DTensor):
        device_mesh = get_device_mesh()
        if placements is None:
            placements = [Replicate()] * device_mesh.ndim
        if global_size is None:
            return DTensor.from_local(input_tensor, device_mesh, placements)
        else:
            return DTensor(input_tensor, device_mesh, placements, size=global_size)
    return input_tensor


def fix_in_gragh_tensor(fx_module: torch.fx.GraphModule, sharding_strategy):

    for node in fx_module.graph.nodes:
        if node.op == 'get_attr':
            with fx_module.graph.inserting_after(node):
                to_dtensor_node = fx_module.graph.call_function(to_dtensor, args=(node, ))

                node.replace_all_uses_with(to_dtensor_node)

                to_dtensor_node.update_arg(0, node)

        create_ops = [
            "torch.ops.aten.scalar_tensor.", "torch.ops.aten.ones.", "torch.ops.aten.empty.",
            "torch.ops.aten.arange.", "torch.ops.aten.zeros."
        ]

        if node.op == 'call_function':
            for op in create_ops:
                if op in _get_qualified_name(node.target):
                    strategy, global_shape = None, None
                    if node.name in sharding_strategy and len(sharding_strategy[node.name]) > 0:
                        strategy = sharding_strategy[node.name][0]
                        # modify the arg of creat ops
                        # (CAUTION!!) device_mesh maybe mock now,
                        # but the local shape need to be consistent of runtime
                        global_shape = node.args[0]
                        local_shape = list(global_shape)
                        device_mesh = get_device_mesh()
                        if isinstance(device_mesh, MockDeviceMesh):
                            logger.warning("maybe wrong shape in MockDeviceMesh for create ops.")
                        for idx, placement in enumerate(strategy):
                            if placement.is_shard():
                                shard_dim = placement.dim
                                split_size, pad_idx = divmod(global_shape[shard_dim],
                                                             device_mesh.size(idx))
                                local_shape[shard_dim] = split_size
                                if device_mesh_rank(device_mesh, idx) < pad_idx:
                                    local_shape[shard_dim] = split_size + 1

                        node.update_arg(0, local_shape)

                    with fx_module.graph.inserting_after(node):
                        to_dtensor_node = fx_module.graph.call_function(to_dtensor,
                                                                        args=(node, strategy,
                                                                              global_shape))

                        node.replace_all_uses_with(to_dtensor_node)

                        to_dtensor_node.update_arg(0, node)

    fx_module.recompile()

    return fx_module


def replace_subsequence_use(node, arg_, redist_node):
    users_node = list(node.users.keys())
    node_next = node.next

    def maybe_replace_node(n: Node) -> Node:
        if n == arg_:
            return redist_node
        else:
            return n

    while node_next.name != "":
        if node_next in users_node:
            new_args = map_arg(node_next.args, maybe_replace_node)
            new_kwargs = map_arg(node_next.kwargs, maybe_replace_node)
            assert isinstance(new_args, tuple)
            assert isinstance(new_kwargs, dict)
            node_next.args = new_args
            node_next.kwargs = new_kwargs
        node_next = node_next.next


def sharding_transform(fx_module: torch.fx.GraphModule, sharding_strategy):
    for node in fx_module.graph.nodes:
        if node.op == 'call_function':
            if node.name in sharding_strategy:
                node_strategy = sharding_strategy[node.name]

                # skip for create ops
                ops_name = _get_qualified_name(node.target)
                if ops_name in [
                        "torch.ops.aten.empty.memory_format", "torch.ops.aten.zeros.default"
                ]:
                    continue

                node_args_flatten = pytree.tree_flatten(node.args)[0]
                invars = [arg for arg in node_args_flatten if isinstance(arg, Node)]
                assert (len(invars) == len(node_strategy))

                for arg_, arg_strategy in zip(invars, node_strategy):
                    with fx_module.graph.inserting_before(node):
                        redist_node = fx_module.graph.call_function(redist_tensor_func,
                                                                    args=(arg_, arg_strategy))

                        # FIXME: may introduce redundancy communication, update_arg for all subsequence use
                        node.replace_input_with(arg_, redist_node)
                        #replace_subsequence_use(node, arg_, redist_node)

    fx_module.graph.on_generate_code(lambda _: insert_spmd_head)

    fx_module.recompile()

    # (fix) %_tensor_constant0 : [#users=1] = get_attr[target=_tensor_constant0]
    fx_module = fix_in_gragh_tensor(fx_module, sharding_strategy)

    return fx_module

easydist/torch/passes/sharding.py
alibaba/easydist
# Copyright (c) 2023, Alibaba Group;
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================

import functools

import torch
from torch.fx.node import _get_qualified_name

from easydist.metashard import view_propagation
from easydist.metashard.annotation import ShardAnnotation, ShardDim
from easydist.metashard.combination import CombinationFunc, ReduceOp

from .device_mesh import device_mesh_world_size

aten = torch.ops.aten


class PresetMetaSPMD:
    op_rules = {}


def register_meta_spmd(op):

    def wrapper(impl):
        op_list = op if isinstance(op, list) else [op]
        for op_i in op_list:
            if isinstance(op_i, str):
                op_name = op_i
            else:
                op_name = _get_qualified_name(op_i)
            PresetMetaSPMD.op_rules[op_name] = impl
        return impl

    return wrapper


def preset_meta_spmd(meta_op, input_args=None):
    if isinstance(meta_op, str):
        args, kwargs = input_args
        return PresetMetaSPMD.op_rules[meta_op](args, kwargs)
    elif meta_op.name in PresetMetaSPMD.op_rules:
        if input_args is None:
            args, kwargs = meta_op.input_args
        else:
            args, kwargs = input_args
        return PresetMetaSPMD.op_rules[meta_op.name](args, kwargs)
    return None, None


@register_meta_spmd([
    "placeholder",
    "get_attr",
])
def meta_spmd_placeholder(args, kwargs):
    input_shape = args.shape
    output_shape = args.shape

    view_ann = view_propagation(input_shape, output_shape, world_size=device_mesh_world_size())
    return view_ann['sharding_ann'], view_ann['combination_ann']


@register_meta_spmd([
    aten.view.default,
    aten._unsafe_view.default,
])
def meta_spmd_view(args, kwargs):
    input_shape = args[0].shape
    output_shape = list(args[1])

    view_ann = view_propagation(input_shape, output_shape, world_size=device_mesh_world_size())
    return view_ann['sharding_ann'], view_ann['combination_ann']


@register_meta_spmd(aten._reshape_alias.default)
def meta_spmd_reshape(args, kwargs):
    input_shape = args[0].shape
    output_shape = list(args[1])
    stride = list(args[2])

    view_ann = view_propagation(input_shape, output_shape, world_size=device_mesh_world_size())
    return view_ann['sharding_ann'], view_ann['combination_ann']


# (TODO) here just replicate all optimizer computation
# need to generate some shard strategy here
@register_meta_spmd([
    aten._foreach_add.Scalar,
    aten._foreach_mul.Scalar,
    aten._foreach_sub.Scalar,
    aten._foreach_div.Scalar,
    aten._foreach_neg.default,
    aten._foreach_reciprocal.default,
    aten._foreach_sqrt.default,
])
def meta_foreach_one_tensor_list(args, kwargs):
    tensor_list = args[0]
    sharding_ann = ShardAnnotation([[[ShardDim.get_noshard_dim()] * len(tensor.shape)
                                     for tensor in tensor_list]])
    return sharding_ann, {}


@register_meta_spmd([
    aten._foreach_add.List,
    aten._foreach_mul.List,
    aten._foreach_sub.List,
    aten._foreach_div.List,
])
def meta_foreach_two_tensor_list(args, kwargs):
    tensor_1_list, tensor_2_list = args[0], args[1]
    sharding_ann = ShardAnnotation([[[ShardDim.get_noshard_dim()] * len(tensor.shape)
                                     for tensor in tensor_1_list],
                                    [[ShardDim.get_noshard_dim()] * len(tensor.shape)
                                     for tensor in tensor_2_list]])
    return sharding_ann, {}


@register_meta_spmd([
    aten._foreach_addcdiv.Scalar,
    aten._foreach_addcmul.Scalar,
])
def meta_foreach_three_tensor_list(args, kwargs):
    tensor_1_list, tensor_2_list, tensor_3_list = args[0], args[1], args[2]
    sharding_ann = ShardAnnotation([[[ShardDim.get_noshard_dim()] * len(tensor.shape)
                                     for tensor in tensor_1_list],
                                    [[ShardDim.get_noshard_dim()] * len(tensor.shape)
                                     for tensor in tensor_2_list],
                                    [[ShardDim.get_noshard_dim()] * len(tensor.shape)
                                     for tensor in tensor_3_list]])
    return sharding_ann, {}


@register_meta_spmd([aten.empty.memory_format, aten.zeros.default])
def meta_create_op(args, kwargs):
    tensor_shape = args[0]

    shard_idx = 1
    combine_ann = {}

    sharding_ann = ShardAnnotation([[ShardDim.get_noshard_dim() for _ in tensor_shape]])

    world_size = device_mesh_world_size()

    for dim_idx, dim_shape in enumerate(tensor_shape):
        if world_size <= dim_shape:
            sharding_ann[0][dim_idx] = ShardDim(shard_idx)
            combine_ann[shard_idx] = functools.partial(CombinationFunc.gather, dim=dim_idx)
            shard_idx += 1

    return sharding_ann, combine_ann


@register_meta_spmd(aten.cat.default)
def meta_cat(args, kwargs):
    tensorlist, dim = args[0:2]
    tensor_shape = tensorlist[0].shape

    shard_idx = 1
    sharding_ann = [ShardDim.get_noshard_dim() for _ in tensor_shape]
    combine_ann = {}

    world_size = device_mesh_world_size()

    for idx in range(len(tensor_shape)):
        if world_size <= tensor_shape[idx] and idx != dim:
            sharding_ann[idx] = ShardDim(shard_idx)
            combine_ann[shard_idx] = functools.partial(CombinationFunc.gather, dim=idx)
            shard_idx += 1

    sharding_ann = ShardAnnotation([sharding_ann for _ in tensorlist])

    return sharding_ann, combine_ann


@register_meta_spmd(aten._log_softmax_backward_data.default)
def meta_log_softmax_backward_data(args, kwargs):
    grad_output, output, dim = args[0:3]
    tensor_shape = grad_output.shape

    shard_idx = 1
    sharding_ann = [ShardDim.get_noshard_dim() for _ in tensor_shape]
    combine_ann = {}

    world_size = device_mesh_world_size()

    for idx in range(len(tensor_shape)):
        if world_size <= tensor_shape[idx] and idx != dim:
            sharding_ann[idx] = ShardDim(shard_idx)
            combine_ann[shard_idx] = functools.partial(CombinationFunc.gather, dim=idx)
            shard_idx += 1

    sharding_ann = ShardAnnotation([sharding_ann] * 2)

    return sharding_ann, combine_ann


@register_meta_spmd(aten.mm.default)
def meta_mm(args, kwargs):
    mat1, mat2 = args[0:2]
    shape_n, shape_m = mat1.shape
    shape_m, shape_p = mat2.shape

    sharding_ann = ShardAnnotation([[ShardDim.get_noshard_dim() for _ in range(2)]
                                    for _ in range(2)])
    shard_idx = 1
    combine_ann = {}

    world_size = device_mesh_world_size()

    if world_size <= shape_n:
        sharding_ann[0][0] = ShardDim(shard_idx)
        combine_ann[shard_idx] = functools.partial(CombinationFunc.gather, dim=0)
        shard_idx += 1

    if world_size <= shape_m:
        sharding_ann[0][1] = ShardDim(shard_idx)
        sharding_ann[1][0] = ShardDim(shard_idx)
        combine_ann[shard_idx] = functools.partial(CombinationFunc.reduce, ops=ReduceOp.SUM)
        shard_idx += 1

    if world_size <= shape_p:
        sharding_ann[1][1] = ShardDim(shard_idx)
        combine_ann[shard_idx] = functools.partial(CombinationFunc.gather, dim=1)
        shard_idx += 1

    return sharding_ann, combine_ann


@register_meta_spmd(aten.bmm.default)
def meta_bmm(args, kwargs):
    mat1, mat2 = args[0:2]
    shape_b_, shape_n, shape_m = mat1.shape
    shape_b, shape_m, shape_p = mat2.shape
    assert shape_b_ == shape_b

    sharding_ann = ShardAnnotation([[ShardDim.get_noshard_dim() for _ in range(3)]
                                    for _ in range(2)])
    shard_idx = 1
    combine_ann = {}

    world_size = device_mesh_world_size()

    if world_size <= shape_b:
        sharding_ann[0][0] = ShardDim(shard_idx)
        sharding_ann[1][0] = ShardDim(shard_idx)
        combine_ann[shard_idx] = functools.partial(CombinationFunc.gather, dim=0)
        shard_idx += 1

    if world_size <= shape_n:
        sharding_ann[0][1] = ShardDim(shard_idx)
        combine_ann[shard_idx] = functools.partial(CombinationFunc.gather, dim=1)
        shard_idx += 1

    if world_size <= shape_m:
        sharding_ann[0][2] = ShardDim(shard_idx)
        sharding_ann[1][1] = ShardDim(shard_idx)
        combine_ann[shard_idx] = functools.partial(CombinationFunc.reduce, ops=ReduceOp.SUM)
        shard_idx += 1

    if world_size <= shape_p:
        sharding_ann[1][2] = ShardDim(shard_idx)
        combine_ann[shard_idx] = functools.partial(CombinationFunc.gather, dim=2)
        shard_idx += 1

    return sharding_ann, combine_ann


@register_meta_spmd(aten.empty_like.default)
def meta_empty_like(args, kwargs):
    tensor_shape = args[0].shape

    world_size = device_mesh_world_size()

    sharding_ann = ShardAnnotation([[ShardDim.get_noshard_dim()] * len(tensor_shape)])
    shard_idx = 1
    combine_ann = {}
    for idx, shape in enumerate(tensor_shape):
        if world_size <= shape:
            sharding_ann[0][idx] = ShardDim(shard_idx)
            combine_ann[shard_idx] = functools.partial(CombinationFunc.gather, dim=idx)
            shard_idx += 1

    return sharding_ann, combine_ann


@register_meta_spmd(aten.native_group_norm.default)
def meta_group_norm(args, kwargs):
    tensor_shape, w_shape, b_shape = args[0].shape, args[1].shape, args[2].shape

    world_size = device_mesh_world_size()
    sharding_ann = ShardAnnotation([[ShardDim.get_noshard_dim()] * len(tensor_shape),
                                    [ShardDim.get_noshard_dim()] * len(w_shape),
                                    [ShardDim.get_noshard_dim()] * len(b_shape)])

    shard_idx = 1
    combine_ann = {}

    if world_size <= tensor_shape[0]:
        sharding_ann[0][0] = ShardDim(shard_idx)
        combine_ann[shard_idx] = [
            functools.partial(CombinationFunc.gather, dim=0),
            functools.partial(CombinationFunc.gather, dim=0),
            functools.partial(CombinationFunc.gather, dim=0)
        ]

    return sharding_ann, combine_ann

easydist/torch/preset_propagation.py
alibaba/easydist
# Copyright (c) 2023, Alibaba Group;
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================

import copy
import logging
import os
from typing import Any, Dict, Iterator, Tuple

import torch
import torch.utils._pytree as pytree
from torch.fx.graph_module import GraphModule
from torch.fx.interpreter import Interpreter
from torch.fx.node import Argument, Target, _get_qualified_name
from torch._subclasses.fake_tensor import FakeTensorMode

from easydist.metashard import MetaOp
from easydist.torch.utils import to_meta
import easydist.config as mdconfig
from .preset_propagation import preset_meta_spmd
from .device_mesh import device_mesh_world_size

logger = logging.getLogger(__name__)


def get_shape_info(node_output):
    if isinstance(node_output, torch.Tensor) or isinstance(node_output,
                                                           torch.nn.parameter.Parameter):
        return {"shape": node_output.shape, "dtype": node_output.dtype}
    else:
        return {}


def meta_exec(self, flat_meta_input=None, enable_fallback=True):

    fake_tensor_mode = FakeTensorMode()

    def wrap_fake(x):
        if isinstance(x, torch.Tensor):
            return fake_tensor_mode.from_tensor(x)
        return x

    if flat_meta_input is None:
        flat_meta_input = self.flat_input_args
    try:
        fake_flat_args = pytree.tree_map(wrap_fake, flat_meta_input)
        fake_args, fake_kwargs = pytree.tree_unflatten(fake_flat_args, self.input_pytree_spec)
        with fake_tensor_mode:
            out = self.func(*fake_args, **fake_kwargs)
    except:
        if enable_fallback:
            out = self.exec()
        else:
            raise RuntimeError(f"failed meta_exec with MetaOp {self}")
    return out


# register the meta_exec function in MetaOp
MetaOp.meta_exec = meta_exec


def to_real(tensor, size=None):

    device = mdconfig.easydist_device

    if size is None:
        size = tensor.size()

    if isinstance(tensor, torch.Tensor) and tensor.is_meta:
        if tensor.dtype == torch.bool:
            return torch.rand(size, dtype=torch.float, device=device) > 0.5
        elif torch.is_floating_point(tensor):
            return torch.rand(size, dtype=tensor.dtype, device=device)
        else:
            return torch.randint(high=1, size=size, dtype=tensor.dtype, device=device)
    return tensor


class EDTorchShardingAnn(Interpreter):

    def __init__(self, module: GraphModule, use_cache=True, garbage_collect_values: bool = True):
        super().__init__(module, garbage_collect_values)
        self.use_cache = use_cache
        self.shape_info = {}
        self.sharding_info = {}

        self.pass_by_ops = ["_operator.getitem"]

    def run(self, *args) -> Any:
        """
        Run `module` via interpretation and return the result.
        Args:
            *args: The arguments to the Module to run, in positional order
            initial_env (Optional[Dict[Node, Any]]): An optional starting environment for execution.
                This is a dict mapping `Node` to any value. This can be used, for example, to
                pre-populate results for certain `Nodes` so as to do only partial evaluation within
                the interpreter.
        Returns:
            Any: The value returned from executing the Module
        """
        self.env = {}

        args = pytree.tree_map(to_meta, args)

        # Positional function args are consumed left-to-right by
        # `placeholder` nodes. Use an iterator to keep track of
        # position and extract those values.
        self.args_iter: Iterator[Any] = iter(args)

        for node in self.module.graph.nodes:

            if node in self.env:
                # Short circuit if we have this value. This could
                # be used, for example, for partial evaluation
                # where the caller has pre-populated `env` with
                # values for a subset of the program.
                continue

            node_output = self.run_node(node)
            self.env[node] = pytree.tree_map(to_meta, node_output)
            self.shape_info[node.name] = pytree.tree_map(get_shape_info, self.env[node])

            if self.garbage_collect_values:
                for to_delete in self.user_to_last_uses.get(node, []):
                    del self.env[to_delete]

            if node.op == 'output':
                # output_val = self.env[node]
                self.env = {}
                return self.sharding_info, self.shape_info

    def placeholder(self, target : 'Target', args : Tuple[Argument, ...], kwargs : Dict[str, Any]) -> Any:
        """
        Execute a ``placeholder`` node. Note that this is stateful:
        ``Interpreter`` maintains an internal iterator over
        arguments passed to ``run`` and this method returns
        next() on that iterator.

        Args:
            target (Target): The call target for this node. See
                `Node <https://pytorch.org/docs/master/fx.html#torch.fx.Node>`__ for
                details on semantics
            args (Tuple): Tuple of positional args for this invocation
            kwargs (Dict): Dict of keyword arguments for this invocation

        Returns:
            Any: The argument value that was retrieved.
        """
        assert isinstance(target, str)

        if "placeholder" not in self.sharding_info:
            self.sharding_info["placeholder"] = {}

        if target.startswith('*'):
            # For a starred parameter e.g. `*args`, retrieve all
            # remaining values from the args list.
            return list(self.args_iter)
        else:
            try:
                output = next(self.args_iter)
                if isinstance(output, torch.Tensor):
                    out_meta_str = str(output)
                    if out_meta_str not in self.sharding_info["placeholder"]:
                        sharding_ann, combination_ann = preset_meta_spmd("placeholder", (output, kwargs))
                        if sharding_ann or combination_ann:
                            self.sharding_info["placeholder"][out_meta_str] = {
                                "sharding_ann": sharding_ann,
                                "combination_ann": combination_ann,
                            }
                return output
            except StopIteration as si:
                if len(args) > 0:
                    return args[0]
                else:
                    raise RuntimeError(f'Expected positional argument for parameter {target}, but one was not passed in!') from si

    def get_attr(self, target : 'Target', args : Tuple[Argument, ...], kwargs : Dict[str, Any]) -> Any:
        """
        Execute a ``get_attr`` node. Will retrieve an attribute
        value from the ``Module`` hierarchy of ``self.module``.

        Args:
            target (Target): The call target for this node. See
                `Node <https://pytorch.org/docs/master/fx.html#torch.fx.Node>`__ for
                details on semantics
            args (Tuple): Tuple of positional args for this invocation
            kwargs (Dict): Dict of keyword arguments for this invocation

        Return:
            Any: The value of the attribute that was retrieved
        """
        assert isinstance(target, str)
        output = self.fetch_attr(target)

        if "get_attr" not in self.sharding_info:
            self.sharding_info["get_attr"] = {}

        if isinstance(output, torch.Tensor):
            out_meta_str = str(output)
            if out_meta_str not in self.sharding_info["get_attr"]:
                sharding_ann, combination_ann = preset_meta_spmd("get_attr", (output, kwargs))
                if sharding_ann or combination_ann:
                    self.sharding_info["get_attr"][out_meta_str] = {
                        "sharding_ann": sharding_ann,
                        "combination_ann": combination_ann,
                    }
        return output

    def call_function(self, target: 'Target', args: Tuple[Argument, ...],
                      kwargs: Dict[str, Any]) -> Any:
        """
        Execute a ``call_function`` node and return the result.
        Args:
            target (Target): The call target for this node. See
                `Node <https://pytorch.org/docs/master/fx.html#torch.fx.Node>`__ for
                details on semantics
            args (Tuple): Tuple of positional args for this invocation
            kwargs (Dict): Dict of keyword arguments for this invocation
        Return
            Any: The value returned by the function invocation
        """
        assert not isinstance(target, str)

        ops_name = _get_qualified_name(target)

        if ops_name in self.pass_by_ops:
            real_out = target(*args, **kwargs)
            return real_out

        args_meta, kwargs_meta = args, kwargs

        op_perf_key = {
            "ops_name": ops_name,
            "args_meta": str(args_meta) + ' | ' + str(kwargs_meta)
        }
        logger.debug(op_perf_key)

        def get_hint_size(args_list, mock_dim_value=1024):
            size_dict = {}
            element_number = 0
            max_dim_value = 0
            for arg_idx in range(len(args_list)):
                if isinstance(args_list[arg_idx], torch.Tensor):
                    tensor_size = args_list[arg_idx].size()
                    size_dict[arg_idx] = tensor_size
                    element_number += tensor_size.numel()
                    if len(args_list[arg_idx].size()) > 0:
                        max_dim_value = max(max_dim_value, max(tensor_size))

            # if element number of all tensor larger then 1024 ** 3,
            hint_size = size_dict
            if mdconfig.use_hint and element_number >= 1024**3:
                hint_size = {}
                for arg_idx in size_dict:
                    hint_size[arg_idx] = torch.Size([
                        mock_dim_value if dim == max_dim_value else dim
                        for dim in size_dict[arg_idx]
                    ])

                logger.info(f"use mock shape for {op_perf_key}, {size_dict} -> {hint_size}")

            return size_dict, hint_size

        # (CAUTION !!) use mock (smaller size) args/kwargs here for OOM issue of giant Op.
        # materialize the input of ops, execute the function and return the result
        def materialize_args_kwargs(args, kwargs):
            flat_args, args_specs = pytree.tree_flatten((args, kwargs))

            ori_size, hint_size = get_hint_size(flat_args)

            def materialize_with_size_dict(size_dict):
                for arg_idx in range(len(flat_args)):
                    if isinstance(flat_args[arg_idx], torch.Tensor):
                        real_tensor = to_real(flat_args[arg_idx], size_dict[arg_idx])

                        # clip for embedding op prevent indexerror in CUDA kernel.
                        if ops_name in [
                                "torch.ops.aten.embedding.default",
                                "easydist.torch.passes.fix_embedding.md_embedding"
                        ]:
                            real_tensor.clip_(max=args[0].shape[0])

                        flat_args[arg_idx] = real_tensor

            materialize_with_size_dict(hint_size)

            # try to run the metaop on the mock size (hint_size)
            # if failed with materialize with the original size
            try:
                meta_op = MetaOp(func=target, input_args=(args, kwargs), name=ops_name)
                meta_op.meta_exec(flat_meta_input=flat_args)
            except:
                materialize_with_size_dict(ori_size)

            return pytree.tree_unflatten(flat_args, args_specs)

        args, kwargs = materialize_args_kwargs(args_meta, kwargs_meta)
        meta_op = MetaOp(func=target,
                         input_args=(args, kwargs),
                         shard_size=device_mesh_world_size(),
                         name=ops_name)
        # user fake tensor here, maybe use shape/dtype info from `make_fx`
        meta_out = meta_op.meta_exec(flat_meta_input=pytree.tree_flatten((args_meta,
                                                                          kwargs_meta))[0])

        # sharding discovery
        if op_perf_key["ops_name"] not in self.sharding_info:
            self.sharding_info[op_perf_key["ops_name"]] = {}

        if not self.use_cache or op_perf_key["args_meta"] not in self.sharding_info[
                op_perf_key["ops_name"]]:
            prompt_annotation = None
            if self.use_cache and len(self.sharding_info[op_perf_key["ops_name"]]) >= 1:
                prompt_annotation = list(
                    self.sharding_info[op_perf_key["ops_name"]].values())[0]["sharding_ann"]

            sharding_ann, combination_ann = preset_meta_spmd(meta_op, (args_meta, kwargs_meta))
            if sharding_ann is None and combination_ann is None:
                sharding_ann, combination_ann = meta_op.sharding_discovery(
                    prompt_annotation=copy.deepcopy(prompt_annotation))

            self.sharding_info[op_perf_key["ops_name"]][op_perf_key["args_meta"]] = {
                "sharding_ann": sharding_ann,
                "combination_ann": combination_ann,
            }

        return pytree.tree_map(to_meta, meta_out)

easydist/torch/sharding_interpreter.py
alibaba/easydist
# Register sharding propagation rules for some ops. These rules are missed in community pytorch.
# refer to https://github.com/pytorch/pytorch/tree/main/torch/distributed/_tensor/ops.

import copy
import functools
import math
import operator
from functools import reduce
from typing import cast

import torch
from torch._meta_registrations import calc_conv_nd_return_shape
from torch.distributed._tensor.op_schema import OpSchema, OutputSharding
from torch.distributed._tensor.ops.common_rules import pointwise_rule
from torch.distributed._tensor.ops.utils import register_prop_rule
from torch.distributed._tensor.placement_types import (DTensorSpec, Replicate, Shard, _Partial)
from easydist.torch.utils import get_dtensor_spec

aten = torch.ops.aten

extra_pointwise_op = [
    aten.silu.default, aten.leaky_relu.default, aten.leaky_relu_backward.default, aten.elu.default,
    aten.elu_backward.default, aten.fill.Scalar
]

for op in extra_pointwise_op:
    register_prop_rule(op)(pointwise_rule)


# from torch/distributed/_spmd/experimental_ops.py on pytorch main branch
# copy it because it is not ready in released pytorch
@register_prop_rule(aten.native_layer_norm.default)  # pyre-ignore
def _prop_native_layer_norm(op_schema: OpSchema) -> OutputSharding:
    input, normalized_shape, weight, bias, eps = op_schema.args_schema
    assert isinstance(input, DTensorSpec)
    assert isinstance(normalized_shape, (tuple, list))
    if weight is not None:
        assert isinstance(weight, DTensorSpec)
        assert all(isinstance(p, Replicate) for p in weight.placements)
    if bias is not None:
        assert isinstance(bias, DTensorSpec)
        assert all(isinstance(p, Replicate) for p in bias.placements)
    # only the left-most (non-normalized) dimensions of the input can be sharded
    batch_ndim = len(input.shape) - len(normalized_shape)
    assert all(
        isinstance(p, Replicate) or (isinstance(p, Shard) and p.dim < batch_ndim, )
        for p in input.placements)
    stats_spec = get_dtensor_spec(
        mesh=input.mesh,
        placements=input.placements,
        shape=torch.Size(s for (i, s) in enumerate(input.shape)),
    )
    return OutputSharding(output_spec=(input, stats_spec, stats_spec))


# from torch/distributed/_spmd/experimental_ops.py on pytorch main branch
# copy it because it is not ready in released pytorch
@register_prop_rule(aten.native_layer_norm_backward.default)  # pyre-ignore
def _prop_native_layer_norm_backward(op_schema: OpSchema) -> OutputSharding:
    (
        grad,
        input,
        normalized_shape,
        result1,
        result2,
        weight,
        bias,
        grad_input_mask,
    ) = op_schema.args_schema
    assert isinstance(grad, DTensorSpec)
    assert isinstance(grad_input_mask, (list, tuple))
    if weight is not None:
        assert isinstance(weight, DTensorSpec)
        assert all(isinstance(s, Replicate) for s in weight.placements)
    if bias is not None:
        assert isinstance(bias, DTensorSpec)
        assert all(isinstance(s, Replicate) for s in bias.placements)

    if all(isinstance(s, Replicate) for s in grad.placements):
        weight_grad = (get_dtensor_spec(
            mesh=weight.mesh,
            placements=[Replicate()] * weight.mesh.ndim,
            shape=weight.shape,
        ) if weight else None)
        bias_grad = (get_dtensor_spec(
            mesh=bias.mesh,
            placements=[Replicate()] * bias.mesh.ndim,
            shape=bias.shape,
        ) if bias else None)
        return OutputSharding(
            # NOTE: type errors below are legit. This is because DTensor currently
            # doesn't support Optional return values. Need to be fixed in DTensor repo.
            output_spec=(
                grad if grad_input_mask[0] else None,
                weight_grad if grad_input_mask[1] else None,
                bias_grad if grad_input_mask[2] else None,
            ), )

    batch_ndim = len(input.shape) - len(normalized_shape)
    assert any(isinstance(s, Shard) and s.dim < batch_ndim
               for s in grad.placements), f"Got {grad.placements}"
    weight_grad = (get_dtensor_spec(
        mesh=weight.mesh,
        placements=[_Partial()] * weight.mesh.ndim,
        shape=weight.shape,
    ) if weight else None)
    bias_grad = (get_dtensor_spec(
        mesh=bias.mesh,
        placements=[_Partial()] * bias.mesh.ndim,
        shape=bias.shape,
    ) if bias else None)
    return OutputSharding(
        # NOTE: type errors below are legit. This is because DTensor currently
        # doesn't support Optional return values. Need to be fixed in DTensor repo.
        output_spec=(
            grad if grad_input_mask[0] else None,
            weight_grad if grad_input_mask[1] else None,
            bias_grad if grad_input_mask[2] else None,
        ), )


@register_prop_rule(aten.convolution.default)
def _prop_convolution_default(op_schema: OpSchema) -> OutputSharding:
    input, weight = op_schema.args_schema[0:2]
    stride, padding, dialation, transposed, output_padding, groups = op_schema.args_schema[-6:]
    assert isinstance(input, DTensorSpec)
    assert isinstance(weight, DTensorSpec)
    output_placements = [Replicate()] * input.mesh.ndim

    if len(output_padding) == 0 or all(v == 0 for v in output_padding):
        output_padding = None
    output_shape = calc_conv_nd_return_shape(input, weight, stride, padding, dialation, transposed,
                                             groups, output_padding)

    input_shape = input.shape
    for idx, s in enumerate(input.placements):
        # if shard on batch, height, witdh
        if isinstance(s, Shard) and s.dim in [0] + list(range(2, len(input_shape))):
            output_placements[idx] = Shard(dim=s.dim)

    for idx, s in enumerate(weight.placements):
        if isinstance(s, Shard) and s.dim == 0:
            output_placements[idx] = Shard(dim=1)
        s_input = input.placements[idx]
        if isinstance(s, Shard) and s.dim == 1 and isinstance(s_input, Shard) and s_input.dim == 1:
            output_placements[idx] = _Partial()

    output_spec = get_dtensor_spec(mesh=input.mesh,
                                   placements=output_placements,
                                   shape=torch.Size(output_shape))

    return OutputSharding(output_spec=output_spec)


@register_prop_rule(aten.convolution_backward.default)
def _prop_convolution_backward_default(op_schema: OpSchema) -> OutputSharding:
    (grad_output, input, weight, bias_sizes, stride, padding, dilation, transposed, output_padding,
     groups, output_mask) = op_schema.args_schema

    weight_placement = [Replicate()] * weight.mesh.ndim
    bias_placement = [Replicate()] * weight.mesh.ndim

    for idx, s in enumerate(input.placements):
        if isinstance(s, Shard) and s.dim == 0:
            weight_placement[idx] = _Partial()
            bias_placement[idx] = _Partial()
        else:
            weight_placement[idx] = weight.placements[idx]
            if isinstance(weight_placement[idx], Shard) and weight_placement[idx].dim == 0:
                bias_placement[idx] = Shard(dim=0)

    weight_grad = get_dtensor_spec(mesh=weight.mesh,
                                   placements=weight_placement,
                                   shape=weight.shape)

    bias_grad = get_dtensor_spec(mesh=weight.mesh,
                                 placements=bias_placement,
                                 shape=torch.Size(bias_sizes))

    return OutputSharding(output_spec=(
        input if output_mask[0] else None,
        weight_grad if output_mask[1] else None,
        bias_grad if output_mask[2] else None,
    ), )


@register_prop_rule(aten.native_batch_norm.default)
def _prop_batch_norm_default(op_schema: OpSchema) -> OutputSharding:
    input, weight, bias = op_schema.args_schema[0:3]
    running_mean, running_var = op_schema.args_schema[3:5]

    return OutputSharding(output_spec=(input, running_mean, running_var))


@register_prop_rule(aten.cudnn_batch_norm.default)
def _prop_batch_norm_default(op_schema: OpSchema) -> OutputSharding:
    input, weight, bias = op_schema.args_schema[0:3]
    running_mean, running_var = op_schema.args_schema[3:5]

    reserve_placement = [Replicate()] * weight.mesh.ndim

    reserve = get_dtensor_spec(mesh=weight.mesh,
                               placements=reserve_placement,
                               shape=torch.Size([0]))

    return OutputSharding(output_spec=(input, running_mean, running_var, reserve), )


@register_prop_rule(aten._native_batch_norm_legit_functional.default)
def _prop_batch_norm_legit_functional_default(op_schema: OpSchema) -> OutputSharding:

    input, weight, bias = op_schema.args_schema[0:3]
    running_mean, running_var = op_schema.args_schema[3:5]

    return OutputSharding(output_spec=(input, running_mean, running_var, running_mean,
                                       running_var), )


@register_prop_rule(aten.native_group_norm.default)
def _prop_native_group_norm(op_schema: OpSchema) -> OutputSharding:

    input, weight, bias, N, C, HxW, group, eps = op_schema.args_schema

    local_shape = input.local_shape
    flatten_inner_shape = functools.reduce((lambda x, y: x * y), local_shape[2:])

    op_schema.args_schema = input, weight, bias, local_shape[0], local_shape[
        1], flatten_inner_shape, group, eps

    assert isinstance(input, DTensorSpec)
    if weight is not None:
        assert isinstance(weight, DTensorSpec)
        assert all(isinstance(p, Replicate) for p in weight.placements)
    if bias is not None:
        assert isinstance(bias, DTensorSpec)
        assert all(isinstance(p, Replicate) for p in bias.placements)
    # only the batch (1st) dimensions of the input can be sharded
    batch_ndim = 1
    assert all(
        isinstance(p, Replicate) or (isinstance(p, Shard) and p.dim < batch_ndim, )
        for p in input.placements)

    stats_placement = [Replicate(), Replicate()]
    for idx, p in enumerate(input.placements):
        if isinstance(p, Shard) and p.dim < batch_ndim:
            stats_placement[idx] = copy.deepcopy(p)
    stats_spec = get_dtensor_spec(
        mesh=input.mesh,
        placements=stats_placement,
        shape=torch.Size((input.shape[0], group)),
    )
    return OutputSharding(output_spec=(input, stats_spec, stats_spec))


@register_prop_rule(aten.native_group_norm_backward.default)
def _prop_native_group_norm_backward(op_schema: OpSchema) -> OutputSharding:
    grad_output, input, mean, rstd, gamma, N, C, HxW, group, output_mask = op_schema.args_schema

    local_shape = input.local_shape

    op_schema.args_schema = grad_output, input, mean, rstd, gamma, local_shape[0], local_shape[
        1], reduce(operator.mul, local_shape[2:]), group, output_mask

    spec = Replicate()
    if isinstance(input.placements[0], Shard):
        spec = _Partial()
    d_gamma = get_dtensor_spec(mesh=gamma.mesh,
                               placements=[spec] * gamma.mesh.ndim,
                               shape=gamma.shape)

    d_bias = get_dtensor_spec(mesh=gamma.mesh,
                              placements=[spec] * gamma.mesh.ndim,
                              shape=gamma.shape)

    return OutputSharding(
        # NOTE: type errors below are legit. This is because DTensor currently
        # doesn't support Optional return values. Need to be fixed in DTensor repo.
        output_spec=(
            input if output_mask[0] else None,
            d_gamma if output_mask[1] else None,
            d_bias if output_mask[2] else None,
        ), )


@register_prop_rule(
    [aten.cudnn_batch_norm_backward.default, aten.native_batch_norm_backward.default])
def _prop_batch_norm_backward_default(op_schema: OpSchema) -> OutputSharding:
    input, grad_output, weight, running_mean, running_var = op_schema.args_schema[0:5]

    weight_placement = [Replicate()] * weight.mesh.ndim

    for idx, s in enumerate(weight.placements):
        grad_place = grad_output.placements[idx]
        if isinstance(s, Replicate) and isinstance(grad_place, Shard) and grad_place.dim == 0:
            weight_placement[idx] = _Partial()
        if isinstance(s, Shard):
            weight_placement[idx] = weight.placements[idx]

    weight_grad = get_dtensor_spec(mesh=weight.mesh,
                                   placements=weight_placement,
                                   shape=weight.shape)
    bias_grad = get_dtensor_spec(mesh=weight.mesh, placements=weight_placement, shape=weight.shape)

    return OutputSharding(
        # NOTE: type errors below are legit. This is because DTensor currently
        # doesn't support Optional return values. Need to be fixed in DTensor repo.
        output_spec=(
            input,
            weight_grad,
            bias_grad,
        ), )


@register_prop_rule(aten.max_pool2d_with_indices.default)
def _prop_max_pool2d_with_indices(op_schema: OpSchema) -> OutputSharding:
    input, kernel, stride = op_schema.args_schema[:3]
    if len(op_schema.args_schema) >= 4:
        padding = op_schema.args_schema[3]
    else:
        padding = [0, 0]

    output_shape = list(input.shape)

    for inner_dim in range(2):
        output_shape[2 + inner_dim] = int(
            math.floor(input.shape[2 + inner_dim] + 2 * padding[inner_dim] -
                       (kernel[inner_dim] - 1) - 1) / stride[inner_dim] + 1)

    output_spec = get_dtensor_spec(mesh=input.mesh,
                                   placements=input.placements,
                                   shape=torch.Size(output_shape))

    return OutputSharding(output_spec=(output_spec, input))


@register_prop_rule([
    aten.max_pool2d_with_indices_backward.default,
    aten.avg_pool3d_backward.default,
])
def _prop_max_pool2d_with_indices_backward(op_schema: OpSchema) -> OutputSharding:
    grad_output, input = op_schema.args_schema[0:2]

    return OutputSharding(output_spec=get_dtensor_spec(
        mesh=grad_output.mesh, placements=grad_output.placements, shape=input.shape))


@register_prop_rule(aten.constant_pad_nd.default)
def _prop_constant_pad_nd(op_schema: OpSchema) -> OutputSharding:
    input, pad, value = op_schema.args_schema

    input_sizes = input.shape
    l_inp = len(input_sizes)

    l_pad = len(pad) // 2
    l_diff = l_inp - l_pad

    new_shape = list(input_sizes[:l_diff])

    for i in range(l_pad):
        pad_idx = len(pad) - ((i + 1) * 2)
        new_dim = input_sizes[l_diff + i] + pad[pad_idx] + pad[pad_idx + 1]
        new_shape.append(new_dim)

    output_spec = get_dtensor_spec(mesh=input.mesh,
                                   placements=input.placements,
                                   shape=torch.Size(new_shape))

    return OutputSharding(output_spec=output_spec)


@register_prop_rule(aten.embedding.default)
def _prop_embedding(op_schema: OpSchema) -> OutputSharding:
    weight_spec, inp_spec = op_schema.args_spec
    if any(placement.is_shard(0) for placement in weight_spec.placements):
        raise NotImplementedError(
            "DTensor does not support row-wise sharded embedding operation yet!")

    out_shape = inp_spec.shape + weight_spec.shape[1:]

    if all(placement.is_replicate()
           for placement in weight_spec.placements) and inp_spec.placements == [Shard(0)]:
        # Embedding table is replicated, input ids are sharded along batch
        # dimension. Output lookups should match input sharding spec in this case.
        return OutputSharding(output_spec=get_dtensor_spec(
            mesh=inp_spec.mesh, placements=inp_spec.placements, shape=out_shape))

    if all(placement.is_replicate() for placement in inp_spec.placements):
        weight_dim_map = weight_spec.dim_map
        output_dim_map = inp_spec.dim_map
        output_dim_map.append(weight_dim_map[1])
        return OutputSharding(output_spec=DTensorSpec.from_dim_map(
            inp_spec.mesh, output_dim_map, [], shape=out_shape))

    return OutputSharding(
        output_spec=None,
        schema_suggestions=[
            OpSchema(
                func_schema=op_schema.func_schema,
                args_schema=(
                    weight_spec,
                    get_dtensor_spec(mesh=inp_spec.mesh,
                                     placements=[Replicate()] * len(inp_spec.placements),
                                     shape=inp_spec.shape),
                ),
                kwargs_schema=op_schema.kwargs_schema,
            )
        ],
    )


# from torch/distributed/_spmd/experimental_ops.py on pytorch main branch
# copy it because it is not ready in released pytorch
@register_prop_rule(aten.select.int)
def _prop_select(op_schema: OpSchema) -> OutputSharding:

    tensor, dim = op_schema.args_schema[:2]
    assert isinstance(tensor, DTensorSpec)
    assert isinstance(dim, int)
    placements = tensor.placements
    assert all(not p.is_shard(dim)
               for p in placements), "DTensor does not support select on sharded dimension."

    # select will remove one dimension, decrement dim of Shard placements by 1
    # if they are larger than dim.
    new_placements = []
    for p in placements:
        # Using isinstance instead of is_shard so that mypy won't complain
        # about accessing dim attribute.
        if isinstance(p, Shard) and p.dim > dim:
            new_placements.append(Shard(p.dim - 1))
        else:
            new_placements.append(p)

    new_shape = tensor.shape[:dim] + tensor.shape[dim + 1:]

    return OutputSharding(
        output_spec=get_dtensor_spec(mesh=tensor.mesh, placements=new_placements, shape=new_shape))


@register_prop_rule(aten.slice_backward.default)
def _prop_slice_backward(op_schema: OpSchema) -> OutputSharding:
    grad, input_sizes = op_schema.args_schema[0:2]

    return OutputSharding(output_spec=get_dtensor_spec(
        mesh=grad.mesh, placements=grad.placements, shape=input_sizes))


@register_prop_rule(  # pyre-ignore
    [
        aten._foreach_add.Scalar,
        aten._foreach_div.Scalar,
        aten._foreach_mul.Scalar,
        aten._foreach_sub.Scalar,
    ])
def _prop__foreach_binop_scalar(op_schema: OpSchema) -> OutputSharding:
    self, scalar = op_schema.args_schema
    assert isinstance(self, list) and all(isinstance(s, DTensorSpec) for s in self)
    assert not isinstance(scalar, list)
    return OutputSharding(output_spec=self)


@register_prop_rule(  # pyre-ignore
    [
        aten._foreach_sub.List,
        aten._foreach_lerp.Scalar,
    ])
def _prop__foreach_trinop_scalar(op_schema: OpSchema) -> OutputSharding:
    self, other = op_schema.args_schema[:2]
    assert isinstance(self, list) and all(isinstance(s, DTensorSpec) for s in self)
    assert isinstance(self, list) and all(isinstance(s, DTensorSpec) for s in other)
    return OutputSharding(output_spec=self)


@register_prop_rule(  # pyre-ignore
    [
        aten._foreach_add.List,
        aten._foreach_div.List,
        aten._foreach_mul.List,
    ])
def _prop__foreach_binop_list(op_schema: OpSchema) -> OutputSharding:
    self, other = op_schema.args_schema[:2]
    scalar = None if len(op_schema.args_schema) < 3 else op_schema.args_schema[2]
    assert isinstance(self, list) and all(
        isinstance(s, DTensorSpec) for s in self), f"Expect a List[DTensorSpec] but got {self}"
    assert isinstance(other, list) and all(
        isinstance(o, DTensorSpec) for o in other), f"Expect a List[DTensorSpec] but got {other}"
    assert len(self) == len(other), ("Two tensor lists must match in length, "
                                     f"but got {len(self)} and {len(other)}")

    # if any(s != o for s, o in zip(self, other)):
    #     # If DTensorSpec for the two operand do not match, suggest using
    #     # self's DTensorSpec. This will trigger allreduce if other is partial
    #     # and self is replicated.
    #     return OutputSharding(
    #         output_spec=None,
    #         schema_suggestions=[
    #             OpSchema(
    #                 func_schema=op_schema.func_schema,
    #                 args_schema=(self, self, scalar) if scalar else (self, self),
    #                 kwargs_schema=op_schema.kwargs_schema,
    #                 is_inplace=op_schema.is_inplace,
    #                 is_out_variant=op_schema.is_out_variant,
    #             )
    #         ],
    #     )
    # else:
    return OutputSharding(output_spec=self)


@register_prop_rule(  # pyre-ignore
    [
        aten._foreach_addcdiv.Scalar,
        aten._foreach_addcmul.Scalar,
    ])
def _prop__foreach_addcop_scalar(op_schema: OpSchema):
    self, tensor1, tensor2 = op_schema.args_schema[:3]
    scalar = None if len(op_schema.args_schema) < 4 else op_schema.args_schema[3]
    assert isinstance(self, list) and all(isinstance(s, DTensorSpec) for s in self)
    assert isinstance(tensor1, list) and all(isinstance(s, DTensorSpec) for s in self)
    assert isinstance(tensor2, list) and all(isinstance(s, DTensorSpec) for s in self)
    if any(s != t1 or s != t2 for s, t1, t2 in zip(self, tensor1, tensor2)):
        # If DTensorSpec for the two operand do not match, suggest using
        # self's DTensorSpec. This will trigger allreduce if other is partial
        # and self is replicated.
        return OutputSharding(
            output_spec=None,
            schema_suggestions=[
                OpSchema(
                    func_schema=op_schema.func_schema,
                    args_schema=(self, self, self, scalar) if scalar else (self, self, self),
                    kwargs_schema=op_schema.kwargs_schema,
                    is_inplace=op_schema.is_inplace,
                    is_out_variant=op_schema.is_out_variant,
                )
            ],
        )
    else:
        return OutputSharding(output_spec=self)


@register_prop_rule(  # pyre-ignore
    [
        aten._foreach_neg.default,
        aten._foreach_reciprocal.default,
        aten._foreach_sqrt.default,
    ])
def _prop__foreach_unaop(op_schema: OpSchema) -> OutputSharding:
    self = op_schema.args_schema[0]
    assert isinstance(self, list) and all(isinstance(s, DTensorSpec) for s in self)
    # FIXME(@mrshenli): for sqrt, this is only mathematically correct for
    # Replicate and Shard tensor.
    return OutputSharding(output_spec=self)


@register_prop_rule(aten.upsample_nearest2d.default)
def _prop_upsample_nearest2d(op_schema: OpSchema) -> OutputSharding:
    input, output_size = op_schema.args_schema[0:2]
    batch, channel = input.shape[:2]
    full_ouput_size = (batch, channel, *output_size)
    return OutputSharding(output_spec=get_dtensor_spec(
        mesh=input.mesh, placements=input.placements, shape=torch.Size(full_ouput_size)))


@register_prop_rule(aten.upsample_nearest2d_backward.default)
def _prop_upsample_nearest2d_backward(op_schema: OpSchema) -> OutputSharding:
    grad_output, output_size, input_size = op_schema.args_schema[0:3]
    return OutputSharding(output_spec=get_dtensor_spec(
        mesh=grad_output.mesh, placements=grad_output.placements, shape=torch.Size(input_size)))


@register_prop_rule(aten.tril.default)
def _prop_tril(op_schema: OpSchema) -> OutputSharding:
    self = op_schema.args_schema[0]

    assert isinstance(self, DTensorSpec)
    assert all(isinstance(p, Replicate) for p in self.placements)

    return OutputSharding(output_spec=self)


# (TODO) merge with _tensor/ops/math_ops.py aten._softmax.default
@register_prop_rule(aten._log_softmax.default)
def _prop_log_softmax_rule(op_schema: OpSchema) -> OutputSharding:
    input_spec, softmax_dim, _ = op_schema.args_schema
    input_spec = cast(DTensorSpec, input_spec)
    softmax_dim = cast(int, softmax_dim)
    dim_map = input_spec.dim_map
    if softmax_dim < len(dim_map) and dim_map[softmax_dim] >= 0:
        raise RuntimeError("Cannot run softmax on sharding dimension!")
    return OutputSharding(input_spec)


# (TODO) merge with _tensor/ops/math_ops.py aten._softmax_backward_data.default
@register_prop_rule(aten._log_softmax_backward_data.default)
def softmax_bwd_rule(op_schema: OpSchema) -> OutputSharding:
    grad_out_spec, out_spec, softmax_dim, _ = op_schema.args_schema
    grad_out_spec = cast(DTensorSpec, grad_out_spec)
    out_spec = cast(DTensorSpec, out_spec)
    softmax_dim = cast(int, softmax_dim)
    grad_out_dim_map = grad_out_spec.dim_map
    out_dim_map = out_spec.dim_map
    if softmax_dim < len(grad_out_dim_map) and (grad_out_dim_map[softmax_dim] >= 0
                                                or out_dim_map[softmax_dim] >= 0):
        raise RuntimeError("Cannot run _softmax_backward_data on sharding dimension!")
    return pointwise_rule(op_schema)


@register_prop_rule(aten.nll_loss_forward.default)  # pyre-ignore
def _prop_nll_loss_forward(op_schema: OpSchema) -> OutputSharding:
    self, target = op_schema.args_schema[:2]
    assert isinstance(self, DTensorSpec)
    assert isinstance(target, DTensorSpec)

    out_placement = [Replicate()] * self.mesh.ndim

    for idx, s in enumerate(self.placements):
        if isinstance(s, Shard) and s.dim == 0:
            out_placement[idx] = _Partial()
        else:
            out_placement[idx] = self.placements[idx]

    return OutputSharding(output_spec=(
        get_dtensor_spec(mesh=self.mesh, placements=out_placement, shape=torch.Size([])),
        get_dtensor_spec(
            mesh=self.mesh, placements=[Replicate()] * self.mesh.ndim, shape=torch.Size([])),
    ))


@register_prop_rule(aten.nll_loss_backward.default)  # pyre-ignore
def _prop_nll_loss_backward(op_schema: OpSchema) -> OutputSharding:
    grad_output, self = op_schema.args_schema[:2]
    assert isinstance(grad_output, DTensorSpec)
    assert isinstance(self, DTensorSpec)
    return OutputSharding(output_spec=self)


@register_prop_rule(aten.avg_pool3d.default)
def _prop_avg_pool3d(op_schema: OpSchema) -> OutputSharding:
    input, kernel_size, stride = op_schema.args_schema[:3]

    padding = [0, 0, 0]
    if len(op_schema.args_schema) > 3:
        padding = op_schema.args_schema[3]

    output_shape = list(input.shape)

    for inner_dim in range(3):
        output_shape[2 + inner_dim] = int(
            math.floor(input.shape[2 + inner_dim] + 2 * padding[inner_dim] -
                       (kernel_size[inner_dim] - 1) - 1) / stride[inner_dim] + 1)

    return OutputSharding(output_spec=get_dtensor_spec(
        mesh=input.mesh, placements=input.placements, shape=torch.Size(output_shape)))


# from torch/distributed/_tensor/ops/random_ops.py on pytorch main branch
# copy it because it is not ready in released pytorch
@register_prop_rule(aten.native_dropout.default)
def dropout_rule(op_schema: OpSchema) -> OutputSharding:
    self_spec = cast(DTensorSpec, op_schema.args_schema[0])

    # NOTE: dropout on a partial tensor should be similar to the case of a replicate tensor
    partial = False
    for placement in self_spec.placements:
        if isinstance(placement, _Partial):
            partial = True
            break

    if partial:
        return OutputSharding(
            None,
            failed_reason="aten.native_dropout.default with _Partial is not supported yet!",
        )
    else:
        return OutputSharding(output_spec=(self_spec, self_spec))


@register_prop_rule([aten.max.dim, aten.min.dim])
def reduce_dim_rule(op_schema: OpSchema) -> OutputSharding:
    self, dim = op_schema.args_schema[:2]
    keepdim = False
    if len(op_schema.args_schema) == 3:
        keepdim = op_schema.args_schema[2]

    self_shape, self_placements = list(self.shape), copy.deepcopy(self.placements)
    if keepdim:
        self_shape[dim] = 1
    else:
        del self_shape[dim]

    out_spec = get_dtensor_spec(mesh=self.mesh, placements=self_placements, shape=self_shape)

    return OutputSharding(output_spec=(out_spec, out_spec))


@register_prop_rule([
    aten.index_put.default,
    aten.index_put_.default,
])
def reduce_dim_rule(op_schema: OpSchema) -> OutputSharding:
    self = op_schema.args_schema[0]
    return OutputSharding(output_spec=self)


@register_prop_rule(aten._scaled_dot_product_efficient_attention.default)
def scaled_dot_product_efficient_attention_rule(op_schema: OpSchema) -> OutputSharding:
    query, key, value = op_schema.args_schema[0:3]

    out_placements = query.placements
    dim_size = len(value.shape)
    for idx, s in enumerate(value.placements):
        if isinstance(s, Shard) and s.dim == dim_size - 1:
            out_placements[idx] = Shard(s.dim)

    out_spec = get_dtensor_spec(mesh=query.mesh,
                                placements=out_placements,
                                shape=torch.Size(list(query.shape[:-1]) + list(value.shape[-1:])))

    state_spec = get_dtensor_spec(mesh=query.mesh,
                                  placements=[Replicate()] * query.mesh.ndim,
                                  shape=torch.Size(list(query.shape[:-2]) + [0]))

    return OutputSharding(output_spec=(out_spec, state_spec))


@register_prop_rule(aten.masked_fill_.Scalar)
def masked_fill_rule(op_schema: OpSchema) -> OutputSharding:
    tensor = op_schema.args_schema[0]
    return OutputSharding(output_spec=tensor)


# from torch/distributed/_tensor/ops/embedding_ops.py on pytorch main branch
# copy it because it is not ready in released pytorch
@register_prop_rule(aten.embedding_dense_backward.default)
def embedding_dense_backward_rules(op_schema: OpSchema) -> OutputSharding:
    grad_output, indices, num_weights = op_schema.args_schema[:3]
    assert isinstance(grad_output, DTensorSpec)
    assert isinstance(indices, DTensorSpec)

    output_shape = (num_weights, ) + grad_output.shape[len(indices.shape):]
    output_placements = []

    for idx in range(len(grad_output.placements)):
        if grad_output.placements[idx] == indices.placements[idx]:
            if isinstance(grad_output.placements[idx], Replicate):
                output_placements.append(Replicate())
            else:
                output_placements.append(_Partial())
        elif grad_output.placements == [_Partial()] and indices.placements == [Replicate()]:
            # The embedding table is replicated and the indices is also replicated
            # (local is a more precise term). This is postional embedding. In this
            # case, gradients for the embmedding table should be Partial.
            output_placements.append(_Partial())
        else:
            raise NotImplementedError("Unsupported embedding dense backward schema:\n"
                                      f"grad_output - {grad_output}\n"
                                      f"indices - {indices}")

    return OutputSharding(output_spec=get_dtensor_spec(
        mesh=indices.mesh, shape=output_shape, placements=output_placements))


# from torch/distributed/_spmd/experimental_ops.py on pytorch main branch
# copy it because it is not ready in released pytorch
@register_prop_rule(aten.stack.default)
def _prop_stack(op_schema: OpSchema) -> OutputSharding:
    tensors = op_schema.args_schema[0]
    dim = 0 if len(op_schema.args_schema) == 1 else cast(int, op_schema.args_schema[1])
    assert (isinstance(tensors, list) and len(tensors) > 0), "expect at least one tensor to stack"
    assert all(isinstance(t, DTensorSpec)
               for t in tensors), f"expect a list of DTensorSpecs, but got {tensors}"
    assert all(t.shape == tensors[0].shape
               for t in tensors), f"expect all tensors to have the same shape, but got {tensors}."
    # TODO: provide schema_suggestions when placements do not match
    assert all(
        t.placements == tensors[0].placements
        for t in tensors), f"expect all tensors to have the same placements, but got {tensors}."
    assert all(
        not p.is_shard(dim)
        for p in tensors[0].placements), "DTensor does not support stack on sharded dimension."

    concat_placements = copy.deepcopy(tensors[0].placements)

    for idx, s in enumerate(concat_placements):
        if isinstance(s, Shard) and s.dim >= dim:
            concat_placements[idx] = Shard(s.dim + 1)

    output_shape = list(tensors[0].shape)
    output_shape.insert(dim, len(tensors))

    return OutputSharding(output_spec=get_dtensor_spec(
        mesh=tensors[0].mesh, shape=output_shape, placements=tensors[0].placements))

# aten._foreach_pow not defined in PyTorch < 2.1.0
if "_foreach_pow" in torch.ops.aten.__dir__():
    @register_prop_rule([aten._foreach_pow.ScalarAndTensor])  # pyre-ignore
    def _prop__foreach_pow_scalar_and_tensor(op_schema: OpSchema):
        scala, exponent = op_schema.args_schema
        assert isinstance(exponent, list) and all(
            isinstance(s, DTensorSpec) for s in exponent
        )
        return OutputSharding(output_spec=exponent)

easydist/torch/spmd_prop_rule.py
alibaba/easydist
# Copyright (c) 2023, Alibaba Group;
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================

import hashlib
from contextlib import contextmanager
from copy import copy
from typing import Any, Dict

import torch
from torch._subclasses.fake_tensor import FakeTensor
from torch.distributed._tensor.placement_types import DTensorSpec
from torch.utils._mode_utils import no_dispatch
import torch.utils._pytree as pytree


def to_meta(node_output):
    if isinstance(node_output, FakeTensor):
        with no_dispatch():
            return torch.zeros_like(node_output, device="meta")
    if type(node_output) is torch.Tensor:
        return node_output.detach().to(device="meta").contiguous()
    elif type(node_output) is torch.nn.parameter.Parameter:
        return node_output.data.detach().to(device="meta").contiguous()
    else:
        return node_output


@contextmanager
def _sharding_ann_env():

    ori_cudnn_enabled = torch.backends.cudnn.enabled
    ori_matmul_allow_tf32 = torch.backends.cuda.matmul.allow_tf32
    ori_cudnn_allow_tf32 = torch.backends.cudnn.allow_tf32

    # (TODO) device failed when convolution_backward in sharding_discovery
    torch.backends.cudnn.enabled = False

    torch.backends.cuda.matmul.allow_tf32 = True
    torch.backends.cudnn.allow_tf32 = True

    try:
        yield
    finally:
        torch.backends.cudnn.enabled = ori_cudnn_enabled
        torch.backends.cuda.matmul.allow_tf32 = ori_matmul_allow_tf32
        torch.backends.cudnn.allow_tf32 = ori_cudnn_allow_tf32


@contextmanager
def _rematerialize_optimizer(
    opt: torch.optim.Optimizer,
    named_states: Dict[str, Any],
    params: Dict[str, torch.nn.Parameter],
):
    assert opt is not None

    # update opt.state with proxy tensors
    orig_states: Dict[str, Any] = copy(opt.state)
    for n in named_states:
        # opt.state's key type is string, but optimizer uses Parameter as keys
        opt.state[params[n]] = named_states[n]  # type: ignore[index]

    # FIXME: support multiple parameter groups
    param_group = opt.param_groups[0]
    orig_params = param_group["params"]
    param_group["params"] = params.values()

    try:
        yield
    finally:
        param_group["params"] = orig_params
        opt.state.update(orig_states)


@contextmanager
def _enable_compile():
    # The return value of torch._utils.is_compiling changes optimizer behavior.
    # We need that function to return True to include optimizer in the graph.
    # See: https://github.com/pytorch/pytorch/blob/a524123c91ab399c9dd6882c1189596dd77e7734/torch/optim/optimizer.py#L41
    def f_true():
        return True

    orig_is_compiling_code = torch._utils.is_compiling.__code__
    torch._utils.is_compiling.__code__ = f_true.__code__
    try:
        yield
    finally:
        torch._utils.is_compiling.__code__ = orig_is_compiling_code


def get_input_signature(*args, **kwargs):

    input_meta_str = pytree.tree_map(to_meta, [args, kwargs])
    return hashlib.sha256(repr(input_meta_str).encode("utf-8")).hexdigest()


def get_dtensor_spec(mesh, placements, shape=None, ndim=None):
    # In PyTorch < 2.1.0: DTensorSpec(mesh, placements, shape=None, ndim=None)
    # In PyTorch >= 2.1.0: DTensorSpec(mesh, placements)
    if "shape" in list(DTensorSpec.__dataclass_fields__.keys()):
        return DTensorSpec(mesh=mesh, placements=placements, shape=shape, ndim=ndim)
    return DTensorSpec(mesh=mesh, placements=tuple(placements))

easydist/torch/utils.py
alibaba/easydist
# Copyright (c) 2023, Alibaba Group;
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================

import functools


def rsetattr(obj, attr, val):
    pre, _, post = attr.rpartition('.')
    return setattr(rgetattr(obj, pre) if pre else obj, post, val)


def rgetattr(obj, attr, *args):

    def _getattr(obj, attr):
        return getattr(obj, attr, *args)

    return functools.reduce(_getattr, [obj] + attr.split('.'))

easydist/utils/__init__.py
alibaba/easydist
# Copyright (c) 2023, Alibaba Group;
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================

ALL_PLATFORM = ["torch", "jax", "tvm"]


class MockDeviceMesh:

    def __init(self):
        pass


class TorchMockDeviceMesh(MockDeviceMesh):

    def __init__(self, *arg, debug_only=False):
        super().__init__()
        self.shape = tuple(arg)
        self.debug_only = debug_only

    def size(self, i):
        return self.shape[i]

    def __str__(self) -> str:
        return f"TorchMockDeviceMesh(shape={self.shape})"

    def __repr__(self) -> str:
        return self.__str__()


class JaxDeviceID:

    def __init__(self, *arg):
        self.shape = tuple(arg)


class JaxMockDeviceMesh(MockDeviceMesh):

    def __init__(self, *arg):
        super().__init__()
        self.device_ids = JaxDeviceID(*arg)


def assert_partial_func_equal(func1, func2):
    assert func1.args == func2.args
    assert func1.keywords == func2.keywords
    assert func1.func.__name__ == func2.func.__name__

easydist/utils/testing.py
alibaba/easydist
# Copyright (c) 2023, Alibaba Group;
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================

import os
import time

from easydist.platform import get_backend
import easydist.config as mdconfig


class EDTimer:

    def __init__(self, func, trials=3, warmup_trials=3, in_ms=True, device=None) -> None:
        self.func = func
        self.warmup_trials = warmup_trials
        self.trials = trials
        self.in_ms = in_ms

        self.device = device
        if self.device == None:
            self.device = mdconfig.easydist_device

        self.backend = get_backend()

    def time(self):
        if self.backend == "jax":
            return self.time_jax()
        elif self.backend == "torch":
            if self.device == "cuda":
                return self.time_torch_cuda()
            elif self.device == "cpu":
                return self.time_cpu()
        return None

    def time_cpu(self):
        for _ in range(self.warmup_trials):
            self.func()

        start_t = time.perf_counter()
        for _ in range(self.trials):
            self.func()

        elapsed_time = time.perf_counter() - start_t
        elapsed_time = elapsed_time / self.trials

        # time elapsed in **milliseconds**
        if self.in_ms:
            return elapsed_time * 1000
        return elapsed_time

    def time_torch_cuda(self):
        import torch

        start_evt = []
        end_evt = []
        for _ in range(0, self.trials):
            start_evt.append(torch.cuda.Event(enable_timing=True))
            end_evt.append(torch.cuda.Event(enable_timing=True))

        for trial_idx in range(0, self.trials + self.warmup_trials):
            evt_idx = trial_idx - self.warmup_trials

            if evt_idx >= 0:
                start_evt[evt_idx].record()

            self.func()

            if evt_idx >= 0:
                end_evt[evt_idx].record()

        torch.cuda.synchronize()
        ops_elapsed_time = 0
        for evt_idx in range(0, self.trials):
            # time elapsed in **milliseconds**
            ops_elapsed_time += start_evt[evt_idx].elapsed_time(end_evt[evt_idx])
        ops_elapsed_time = ops_elapsed_time / self.trials

        if self.in_ms:
            return ops_elapsed_time
        return ops_elapsed_time / 1000

    def time_jax(self):
        import jax
        for _ in range(self.warmup_trials):
            self.func()
            (jax.device_put(0.) + 0).block_until_ready()

        start_t = time.perf_counter()
        for _ in range(self.trials):
            self.func()
            (jax.device_put(0.) + 0).block_until_ready()

        elapsed_time = time.perf_counter() - start_t
        elapsed_time = elapsed_time / self.trials

        # time elapsed in **milliseconds**
        if self.in_ms:
            return elapsed_time * 1000
        return elapsed_time

easydist/utils/timer.py
alibaba/easydist
# Copyright (c) 2023, Alibaba Group;
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================

import re

VERSION = "0.1.0"


def is_release_version():
    return bool(re.match(r"^\d+\.\d+\.\d+$", VERSION))

easydist/version.py
alibaba/easydist
# mpirun -np 2 python ./examples/jax/simple_function.py

import logging

import jax
import jax.numpy as jnp

from easydist import easydist_setup, mdconfig
from easydist.jax.api import easydist_compile


@easydist_compile()
def foo_func(x, y):
    tanh = jnp.tanh(x)
    return jnp.exp(tanh) @ y + tanh


def main():
    mdconfig.log_level = logging.INFO
    easydist_setup(backend="jax", device="cuda")

    key = jax.random.PRNGKey(0)
    key, subkey = jax.random.split(key)
    randn_x = jax.random.normal(key, (10, 10))
    randn_y = jax.random.normal(subkey, (10, 10))

    jax_out = foo_func.original_func(randn_x, randn_y)
    md_out = foo_func(randn_x, randn_y)

    if not jax.numpy.allclose(jax_out, md_out):
        raise RuntimeError("simlpe function test failed!!")

    print("simlpe function example pass.")


if __name__ == '__main__':
    main()

examples/jax/simple_function.py
alibaba/easydist
# mpirun -np 2 python ./examples/jax/simple_model.py --mode inference

import argparse
import logging

import jax
import optax
from flax import linen as nn
from flax.training import train_state
from jax import random

from easydist import easydist_setup, mdconfig
from easydist.jax.api import easydist_compile


class Foo(nn.Module):
    """A simple CNN model."""

    @nn.compact
    def __call__(self, x):
        x = nn.LayerNorm()(x)
        x = nn.Dense(features=6)(x)
        x = nn.relu(x)
        return x


def inference_example(module, params, input):

    @easydist_compile()
    def inference_step(params, input):
        out = module.apply({'params': params}, input)
        return out

    jax_out = inference_step.original_func(params, input)
    md_out = inference_step(params, input)

    if not jax.numpy.allclose(jax_out, md_out):
        raise RuntimeError("simlpe model test failed!!")

    print("simlpe model inference example pass.")


def train_example(module, params, input):

    tx = optax.adam(learning_rate=0.01)

    state = train_state.TrainState.create(apply_fn=module.apply, params=params, tx=tx)

    @easydist_compile()
    def train_step(state, batch):
        """Train for a single step."""

        def loss_fn(params):
            logits = state.apply_fn({'params': params}, batch)
            loss = logits.mean()
            return loss

        grad_fn = jax.grad(loss_fn)
        grads = grad_fn(state.params)
        state = state.apply_gradients(grads=grads)
        return state

    jax_state = train_step.original_func(state, input)
    md_state = train_step(state, input)

    if not jax.tree_util.tree_all(jax.tree_map(jax.numpy.allclose, jax_state, md_state)):
        raise RuntimeError("simlpe model train example failed!!")

    print("simlpe model train example pass.")


def main():
    parser = argparse.ArgumentParser(description="Simple example of parallelize model.")

    parser.add_argument("--mode",
                        type=str,
                        default=None,
                        choices=["train", "inference"],
                        required=True)

    args = parser.parse_args()

    mdconfig.log_level = logging.INFO
    easydist_setup(backend="jax", device="cuda")

    model = Foo()

    root_key = jax.random.PRNGKey(seed=0)
    main_key, params_key = jax.random.split(key=root_key, num=2)
    rand_input = random.normal(main_key, (4, 6))
    variables = model.init(params_key, rand_input)
    params = variables['params']

    if args.mode == "train":
        train_example(model, params, rand_input)
    if args.mode == "inference":
        inference_example(model, params, rand_input)


if __name__ == "__main__":
    main()

examples/jax/simple_model.py
alibaba/easydist
# code modified from https://github.com/pytorch/tutorials/blob/main/beginner_source/blitz/cifar10_tutorial.py

import logging
import os
import random

import torch
import torch.nn as nn
import torch.optim as optim
import torchvision
from torchvision.models import resnet18
import torchvision.transforms as transforms

from easydist import easydist_setup, mdconfig
from easydist.torch.experimental.api import easydist_compile

random.seed(42)
torch.manual_seed(42)


@easydist_compile
def train_step(net, optimizer, inputs, labels):

    criterion = nn.CrossEntropyLoss()

    outputs = net(inputs)
    loss = criterion(outputs, labels)
    loss.backward()
    optimizer.step()

    optimizer.zero_grad()

    return loss


def main():

    # setting up easydist and torch.distributed
    mdconfig.log_level = logging.INFO
    easydist_setup(backend="torch", device="cuda")

    torch.distributed.init_process_group(backend="nccl")
    local_rank = int(os.environ["LOCAL_RANK"])
    torch.cuda.set_device(local_rank)

    transform = transforms.Compose(
        [transforms.ToTensor(),
         transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])

    batch_size = 128

    trainset = torchvision.datasets.CIFAR10(root='./data',
                                            train=True,
                                            download=True,
                                            transform=transform)
    trainloader = torch.utils.data.DataLoader(trainset,
                                              batch_size=batch_size,
                                              shuffle=True,
                                              drop_last=True,
                                              num_workers=2)

    net = resnet18().cuda()

    optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9)

    for epoch in range(2):  # loop over the dataset multiple times

        running_loss = 0.0
        for i, data in enumerate(trainloader, 0):
            # get the inputs; data is a list of [inputs, labels]
            inputs, labels = data[0].cuda(), data[1].cuda()

            loss = train_step(net, optimizer, inputs, labels)

            # print statistics
            running_loss += loss.item()
            if i % 100 == 99:  # print every 100 mini-batches
                print(f'[{epoch + 1}, {i + 1:5d}] loss: {running_loss / 2000:.3f}')
                running_loss = 0.0

    print('Finished Training')


if __name__ == '__main__':
    main()

examples/torch/cifar10.py
alibaba/easydist
# code modified from https://github.com/Diego999/pyGAT

import dgl


def load_data(dataset="cora"):

    if dataset == "cora":
        dataset = dgl.data.CoraGraphDataset()
        adj = dataset[0].adj().to_dense()

        features = dataset[0].ndata["feat"]
        labels = dataset[0].ndata["label"]

        train_mask = dataset[0].ndata["train_mask"]
        val_mask = dataset[0].ndata["val_mask"]
        test_mask = dataset[0].ndata["test_mask"]

        idx_train = (train_mask == True).nonzero().flatten()
        idx_val = (val_mask == True).nonzero().flatten()
        idx_test = (test_mask == True).nonzero().flatten()

        return adj, features, labels, idx_train, idx_val, idx_test, train_mask

    elif dataset == "wiki-cs":
        dataset = dgl.data.WikiCSDataset()
        adj = dataset[0].adj().to_dense()

        features = dataset[0].ndata["feat"]
        labels = dataset[0].ndata["label"]

        train_mask = dataset[0].ndata["train_mask"][:, 0]
        val_mask = dataset[0].ndata["val_mask"][:, 0]
        test_mask = dataset[0].ndata["test_mask"]

        idx_train = (train_mask == True).nonzero().flatten()
        idx_val = (val_mask == True).nonzero().flatten()
        idx_test = (test_mask == True).nonzero().flatten()

        return adj, features, labels, idx_train, idx_val, idx_test, train_mask

examples/torch/gnn/data.py
alibaba/easydist
# code modified from https://github.com/Diego999/pyGAT

import torch
import torch.nn as nn
import torch.nn.functional as F


class GATLayer(nn.Module):

    def __init__(self, in_features, out_features, dropout, alpha, concat=True):
        super(GATLayer, self).__init__()
        self.dropout = dropout
        self.in_features = in_features
        self.out_features = out_features
        self.alpha = alpha
        self.concat = concat

        self.linear_w = nn.Linear(in_features=in_features, out_features=out_features, bias=None)
        nn.init.xavier_uniform_(self.linear_w.weight.data, gain=1.414)

        self.linear_a_1 = nn.Linear(in_features=out_features, out_features=1, bias=None)
        self.linear_a_2 = nn.Linear(in_features=out_features, out_features=1, bias=None)
        nn.init.xavier_uniform_(self.linear_a_1.weight.data, gain=1.414)
        nn.init.xavier_uniform_(self.linear_a_2.weight.data, gain=1.414)

        self.leakyrelu = nn.LeakyReLU(self.alpha)

        self.softmax = nn.Softmax(dim=-1)
        self.dropout = nn.Dropout(self.dropout)

        self.register_buffer("zero_vec", None)

    def forward(self, h, adj):
        wh = self.linear_w(h)
        wh1 = self.linear_a_1(wh)
        wh2 = self.linear_a_2(wh)

        e = self.leakyrelu(wh1 + wh2.T)

        if self.zero_vec is None:
            self.zero_vec = -10e10 * torch.ones_like(e)
        attention = torch.where(adj > 0, e, self.zero_vec)
        attention = self.dropout(self.softmax(attention))

        h_new = torch.matmul(attention, wh)

        if self.concat:
            return F.elu(h_new)
        else:
            return h_new


class GAT(nn.Module):

    def __init__(self, nfeat, nhid, nclass, dropout, alpha, nheads):
        """Dense version of GAT."""
        super(GAT, self).__init__()
        self.dropout = dropout

        self.attentions = nn.ModuleList([
            GATLayer(nfeat, nhid, dropout=dropout, alpha=alpha, concat=True) for _ in range(nheads)
        ])

        self.last_layer = GATLayer(nhid * nheads,
                                   nclass,
                                   dropout=dropout,
                                   alpha=alpha,
                                   concat=False)

        self.dropout_1 = nn.Dropout(self.dropout)
        self.dropout_2 = nn.Dropout(self.dropout)

    def forward(self, x, adj):
        x = self.dropout_1(x)

        multi_head_att = []
        for att in self.attentions:
            multi_head_att.append(att(x, adj))

        x = torch.cat(multi_head_att, dim=1)
        x = self.dropout_2(x)
        x = F.elu(self.last_layer(x, adj))

        return F.log_softmax(x, dim=1)

examples/torch/gnn/gat.py
alibaba/easydist
# code modified from https://github.com/Diego999/pyGAT

import argparse
import logging
import os
import sys
import random
import time

import rich
import torch
import torch.nn.functional as F

from easydist import easydist_setup, mdconfig
from easydist.torch.experimental.api import easydist_compile

sys.path.append(os.path.abspath(__file__))
from gat import GAT
from data import load_data

random.seed(42)
torch.manual_seed(42)


def accuracy(output, labels):
    preds = output.max(1)[1].type_as(labels)
    correct = preds.eq(labels).double()
    correct = correct.sum()
    return correct / len(labels)


def get_args():
    parser = argparse.ArgumentParser()

    parser.add_argument('--hidden', type=int, default=4096)
    parser.add_argument('--nb_heads', type=int, default=4)
    parser.add_argument('--dropout', type=float, default=0.5)
    parser.add_argument('--alpha', type=float, default=0.2)

    parser.add_argument('--lr', type=float, default=0.0005)
    parser.add_argument('--weight_decay', type=float, default=5e-4)

    parser.add_argument('--dataset', type=str, choices=["wiki-cs", "cora"], default="wiki-cs")
    parser.add_argument('--epochs', type=int, default=800)

    args = parser.parse_args()
    rich.print(f"Training config: {args}")

    return args


def main():

    # setting up easydist and torch.distributed
    mdconfig.log_level = logging.INFO
    easydist_setup(backend="torch", device="cuda")

    torch.distributed.init_process_group(backend="nccl")
    local_rank = int(os.environ["LOCAL_RANK"])
    torch.cuda.set_device(local_rank)

    args = get_args()

    adj, features, labels, idx_train, idx_val, _, train_mask = load_data(args.dataset)

    model = GAT(nfeat=features.shape[1],
                nhid=args.hidden,
                nclass=int(labels.max()) + 1,
                dropout=args.dropout,
                nheads=args.nb_heads,
                alpha=args.alpha).cuda()

    adj, features, labels, idx_train = adj.cuda(), features.cuda(), labels.cuda(), idx_train.cuda()
    train_mask = train_mask.cuda()

    optimizer = torch.optim.Adam(model.parameters(),
                                 lr=args.lr,
                                 weight_decay=args.weight_decay,
                                 foreach=True,
                                 capturable=True)

    loss_scale = len(idx_train) / len(labels)

    @easydist_compile()
    def train_step(model, optimizer, adj, features, labels, mask):
        output = model(features, adj)
        loss_train = F.nll_loss(output * mask[:, None], labels * mask) / loss_scale
        loss_train.backward()
        optimizer.step()
        optimizer.zero_grad()

        return output, loss_train

    for epoch in range(args.epochs):

        start_t = time.perf_counter()

        output, loss_train = train_step(model, optimizer, adj, features, labels, train_mask)

        acc_train = accuracy(output[idx_train], labels[idx_train])

        loss_val = F.nll_loss(output[idx_val], labels[idx_val])
        acc_val = accuracy(output[idx_val], labels[idx_val])

        rich.print('epoch={:04d} | '.format(epoch + 1),
                   'loss_train={:.6f}'.format(loss_train.data.item()),
                   'acc_train={:.6f}'.format(acc_train.data.item()),
                   'loss_val={:.6f}'.format(loss_val.data.item()),
                   'acc_val={:.6f}'.format(acc_val.data.item()),
                   'time={:.3f}'.format(time.perf_counter() - start_t))


if __name__ == "__main__":
    main()

examples/torch/gnn/train.py
alibaba/easydist
import os
import logging

import torch

from easydist import easydist_setup, mdconfig
from easydist.torch import set_device_mesh
from easydist.torch.experimental.api import easydist_compile
from easydist.utils.testing import TorchMockDeviceMesh


@easydist_compile()
def foo_func(x, y):
    tanh = torch.tanh(x)
    return torch.mm(torch.exp(tanh), y) + tanh


def main():
    mdconfig.log_level = logging.INFO
    easydist_setup(backend="torch", device="cuda", allow_tf32=False)

    torch.distributed.init_process_group(backend="nccl")
    local_rank = int(os.environ["LOCAL_RANK"])
    torch.cuda.set_device(local_rank)

    mesh = TorchMockDeviceMesh(1, 4, debug_only=True)
    set_device_mesh(mesh)

    randn_x = torch.randn(10, 10, requires_grad=True).cuda()
    randn_y = torch.randn(10, 10, requires_grad=True).cuda()
    torch.distributed.broadcast(randn_x, src=0)
    torch.distributed.broadcast(randn_y, src=0)

    torch_out = foo_func.original_func(randn_x, randn_y)
    md_out = foo_func(randn_x, randn_y)

    if not torch.allclose(torch_out, md_out):
        raise RuntimeError("simlpe function test failed!!")

    print("simlpe function example pass.")


if __name__ == '__main__':
    main()

examples/torch/simple_function.py
alibaba/easydist
import argparse
import copy
import logging
import os
from contextlib import nullcontext

import torch
from torch._subclasses.fake_tensor import FakeTensorMode
from torch.distributed.distributed_c10d import _get_default_group
from torch.distributed.utils import _sync_module_states
from torch.utils.checkpoint import checkpoint

from easydist import easydist_setup, mdconfig
from easydist.torch.experimental.api import easydist_compile
from easydist.torch import set_device_mesh
from easydist.utils.testing import TorchMockDeviceMesh


def broadcast_module(model):
    _sync_module_states(model,
                        _get_default_group(),
                        broadcast_bucket_size=int(250 * 1024 * 1024),
                        src=0,
                        params_and_buffers_to_ignore=set())

    return model


class Foo(torch.nn.Module):

    def __init__(self, enable_checkpoint=False):
        super().__init__()
        self.enable_checkpoint = enable_checkpoint
        self.norm = torch.nn.LayerNorm(5)
        self.linear = torch.nn.Linear(5, 5)

    def forward(self, x):
        x = self.norm(x)
        if self.enable_checkpoint:
            x = checkpoint(self.linear, x, preserve_rng_state=False)
        else:
            x = self.linear(x)
        return x.relu()


@easydist_compile(tracing_mode="fake")
@torch.inference_mode()
def inference_step(model, input):
    out = model(input)
    return out


def inference_example(fake_init=True, cpu_init_helper=False):
    fake_mode = FakeTensorMode()

    # (NOTE) initialize cuda context first see https://github.com/pytorch/pytorch/issues/92627
    torch.ones(1).cuda()
    with torch.device('cuda'), fake_mode if fake_init else nullcontext():
        model = Foo()
        randn_input = torch.randn(3, 5)

    # broadcast the parameter and input
    model = broadcast_module(model)
    torch.distributed.broadcast(randn_input, src=0)

    torch_out = inference_step.original_func(model, randn_input)

    mock_mesh = TorchMockDeviceMesh(1, 2, debug_only=True)
    set_device_mesh(mock_mesh)

    if fake_init:
        randn_input = torch.randn(3, 5).cuda()
        torch.distributed.broadcast(randn_input, src=0)

    if cpu_init_helper:
        module = Foo().cuda()
        module = broadcast_module(module)

        inference_step.register_cpu_module(copy.deepcopy(module).to(device="cpu"))

        torch_out = inference_step.original_func(module, randn_input)

    md_out = inference_step(model, randn_input)

    if fake_init and not cpu_init_helper:
        assert md_out.shape == torch_out.shape, "shape mismatch"
    else:
        assert torch.allclose(torch_out, md_out), "simple model inference test failed."

    print("simple model inference example pass.")


# when using cuda_graph, because of the warm-up and cuda graph capture,
# the result of the first step is equivalent to the original result of the third step
@easydist_compile(tracing_mode="fake", cuda_graph=False)
def train_step(input, model, opt):
    out = model(input).mean()
    out.backward()
    opt.step()
    opt.zero_grad(True)

    return out


def train_example(fake_init=True, enable_checkpoint=False, cpu_init_helper=False):
    fake_mode = FakeTensorMode()

    # (NOTE) initialize cuda context first see https://github.com/pytorch/pytorch/issues/92627
    torch.ones(1).cuda()
    with torch.device('cuda'), fake_mode if fake_init else nullcontext():
        model = Foo(enable_checkpoint)
        randn_input = torch.randn(16, 5)

        # broadcast the parameter and input
        model = broadcast_module(model)
        torch.distributed.broadcast(randn_input, src=0)

        opt = torch.optim.Adam(model.parameters(), lr=0.001, foreach=True, capturable=True)

        model_2 = copy.deepcopy(model)
        opt_2 = torch.optim.Adam(model_2.parameters(), lr=0.001, foreach=True, capturable=True)

        torch_step_1_result = train_step.original_func(randn_input, model, opt)
        torch_step_2_result = train_step.original_func(randn_input, model, opt)

    # trace train step func
    mock_mesh = TorchMockDeviceMesh(1, 2, debug_only=True)
    set_device_mesh(mock_mesh)

    # need real input for compiled func
    if fake_init:
        randn_input = torch.randn(3, 5).cuda()
        torch.distributed.broadcast(randn_input, src=0)

    if cpu_init_helper:
        module = Foo(enable_checkpoint).cuda()
        module = broadcast_module(module)
        train_step.register_cpu_module(copy.deepcopy(module).to("cuda"))

        opt = torch.optim.Adam(module.parameters(), lr=0.001, foreach=True, capturable=True)

        torch_step_1_result = train_step.original_func(randn_input, module, opt)
        torch_step_2_result = train_step.original_func(randn_input, module, opt)

    md_step_1_result = train_step(randn_input, model_2, opt_2)
    md_step_2_result = train_step(randn_input, model_2, opt_2)

    if fake_init and not cpu_init_helper:
        assert torch_step_1_result.shape == md_step_1_result.shape, "shape mismatch"
        assert torch_step_2_result.shape == md_step_2_result.shape, "shape mismatch"
    else:
        assert torch.allclose(torch_step_1_result,
                              md_step_1_result), "simple model training test failed."
        assert torch.allclose(torch_step_2_result,
                              md_step_2_result), "simple model training test failed."

    print("simple model training example pass.")


def main():
    parser = argparse.ArgumentParser(description="Simple example of parallelize model.")

    parser.add_argument("--mode",
                        type=str,
                        default=None,
                        choices=["train", "inference"],
                        required=True)
    parser.add_argument("--fake-init", action="store_true")
    parser.add_argument("--checkpoint", action="store_true")
    parser.add_argument("--cpu-init-helper", action="store_true")

    args = parser.parse_args()

    # setting up easydist and torch.distributed
    mdconfig.log_level = logging.INFO
    easydist_setup(backend="torch", device="cuda", allow_tf32=False)

    torch.distributed.init_process_group(backend="nccl")
    local_rank = int(os.environ["LOCAL_RANK"])
    torch.cuda.set_device(local_rank)

    if args.mode == "train":
        train_example(fake_init=args.fake_init,
                      enable_checkpoint=args.checkpoint,
                      cpu_init_helper=args.cpu_init_helper)
    if args.mode == "inference":
        inference_example(fake_init=args.fake_init, cpu_init_helper=args.cpu_init_helper)


if __name__ == "__main__":
    main()

examples/torch/simple_model.py
alibaba/easydist
import copy
import functools
import logging
import os

import diffusers
import torch
import torch.utils._pytree as pytree
from diffusers import EulerDiscreteScheduler, StableDiffusionPipeline

from easydist import easydist_setup, mdconfig
from easydist.torch.experimental.api import easydist_compile

pytree._register_pytree_node(
    diffusers.models.unet_2d_condition.UNet2DConditionOutput, lambda x: ([x.sample], None),
    lambda values, _: diffusers.models.unet_2d_condition.UNet2DConditionOutput(values[0]))


def main():
    # setting up easydist and torch.distributed
    mdconfig.log_level = logging.INFO
    easydist_setup(backend="torch", device="cuda")

    torch.distributed.init_process_group(backend="nccl")
    local_rank = int(os.environ["LOCAL_RANK"])
    torch.cuda.set_device(local_rank)
    torch.manual_seed(42)

    model_id = "stabilityai/stable-diffusion-2"

    # Use the Euler scheduler here instead
    scheduler = EulerDiscreteScheduler.from_pretrained(model_id, subfolder="scheduler")
    pipe = StableDiffusionPipeline.from_pretrained(model_id,
                                                   scheduler=scheduler,
                                                   torch_dtype=torch.float16)
    pipe = pipe.to("cuda")

    @easydist_compile(use_hint=True)
    @torch.inference_mode()
    def sharded_unet(model, *args, **kwargs):
        return model(*args, **kwargs)

    pipe.unet.forward = functools.partial(sharded_unet, copy.copy(pipe.unet))

    prompt = "a photo of Pride and Prejudice"
    image = pipe(prompt, width=1024, height=1024).images[0]
    image.save("pride_and_prejudice.png")


if __name__ == '__main__':
    main()

examples/torch/stable_diffusion.py
alibaba/easydist
from typing import List

import torch
import torch._dynamo


def foo(x, y):
    a = torch.sin(x)
    b = torch.cos(x)
    return a + b


def custom_backend(gm: torch.fx.GraphModule, example_inputs: List[torch.Tensor]):
    print("custom backend called with FX graph:")
    print(gm.graph)
    return gm.forward


model_exp = torch._dynamo.export(foo,
                                 torch.randn(10, 10),
                                 torch.randn(10, 10),
                                 aten_graph=True,
                                 tracing_mode="fake")
model_exp[0].print_readable()
# opt_foo1 = torch.compile(foo, backend=custom_backend, fullgraph=True)
a, b = torch.randn(10, 10), torch.randn(10, 10)
print(foo(a, b))

print(model_exp[0](a, b))

examples/torch/test_dynamo_export.py
alibaba/easydist
import logging

import rich
import tvm
import tvm.testing
from tvm import te
import numpy as np

import easydist as md

logging.basicConfig(format='%(asctime)s - %(levelname)s - %(message)s',
                    datefmt='%m/%d %H:%M:%S',
                    level=logging.DEBUG)

md.platform.init_backend("tvm")

tgt = tvm.target.Target(target="llvm", host="llvm")

n = te.var("n")
A = te.placeholder((n, ), name="A")
B = te.placeholder((n, ), name="B")
C = te.compute(A.shape, lambda i: A[i] + B[i], name="C")

s = te.create_schedule(C.op)

fadd = tvm.build(s, [A, B, C], tgt, name="myadd")


def fadd_wrapped(a, b):
    c = md.platform.zeros_like(a)
    assert a.shape == b.shape
    fadd(a, b, c)
    return c


dev = tvm.device(tgt.kind.name, 0)

n = 1024
a = tvm.nd.array(np.random.uniform(size=n).astype(A.dtype), dev)
b = tvm.nd.array(np.random.uniform(size=n).astype(B.dtype), dev)

meta_op = md.metashard.MetaOp(fadd_wrapped, ((a, b), {}))
sharding_annotion, combination_ann = meta_op.sharding_discovery()

rich.print(sharding_annotion)
rich.print(combination_ann)

c = fadd_wrapped(a, b)
print(c)
print(a.numpy() + b.numpy())

examples/tvm/test_simple.py
alibaba/easydist
# Copyright (c) 2023, Alibaba Group;
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================

import os
import setuptools
from importlib.machinery import SourceFileLoader

version = (
    SourceFileLoader("easydist.version", os.path.join(
        "easydist", "version.py")).load_module().VERSION
)


def is_comment_or_empty(line):
    stripped = line.strip()
    return stripped == "" or stripped.startswith("#")


def remove_comments_and_empty_lines(lines):
    return [line for line in lines if not is_comment_or_empty(line)]


def get_core_requirements():
    with open(os.path.join("requirements", "core-requirements.txt")) as f:
        core_requirements = remove_comments_and_empty_lines(
            f.read().splitlines())
    return core_requirements


def get_long_description():
    with open("README.md", "r") as fh:
        long_description = fh.read()
    return long_description


setuptools.setup(
    name="pai-easydist",
    version=version,
    author="Shenggan Cheng",
    author_email="shenggan.c@u.nus.edu",
    description="Efficient Automatic Training System for Super-Large Models",
    long_description=get_long_description(),
    long_description_content_type="text/markdown",
    url="https://github.com/alibaba/easydist",
    packages=setuptools.find_packages(),
    install_requires=get_core_requirements(),
    extras_require={
        "torch": [
            "torch",
            "torchvision",
        ],
        "jax": [
            "jax[cuda11_pip]",
            "flax",
        ]
    },
)

setup.py
alibaba/easydist
# Copyright (c) 2023, Alibaba Group;
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================

import numpy
import pytest

from easydist import platform
from easydist.metashard.combination import CombinationFunc
from easydist.utils.testing import ALL_PLATFORM
from easydist import easydist_setup


@pytest.mark.parametrize("backend", ALL_PLATFORM)
def test_gather(backend):
    easydist_setup(backend)
    shard_tensor = [platform.from_numpy(numpy.ones((3, 4)))] * 4
    global_tensor_1 = platform.from_numpy(numpy.ones((12, 4)))
    gather_dim1 = CombinationFunc.gather(shard_tensor, dim=0)

    assert platform.allclose(global_tensor_1, gather_dim1)

    global_tensor_2 = platform.from_numpy(numpy.ones((3, 16)))
    gather_dim2 = CombinationFunc.gather(shard_tensor, dim=1)

    assert platform.allclose(global_tensor_2, gather_dim2)


@pytest.mark.parametrize("backend", ALL_PLATFORM)
def test_gather_halo(backend):
    easydist_setup(backend)
    shard_tensor = [platform.from_numpy(numpy.array([1, 1, 1]))] * 3
    global_tensor_1 = platform.from_numpy(numpy.array([1, 1, 2, 1, 2, 1, 1]))
    gather_halo_1 = CombinationFunc.gather(shard_tensor, dim=0, halowidth=1)

    assert platform.allclose(global_tensor_1, gather_halo_1)

    global_tensor_2 = platform.from_numpy(numpy.array([1, 1, 1, 1, 1]))
    gather_halo_2 = CombinationFunc.gather(shard_tensor, dim=0, halowidth=-1)

    assert platform.allclose(global_tensor_2, gather_halo_2)


@pytest.mark.parametrize("backend", ALL_PLATFORM)
def test_gather_chunk(backend):
    easydist_setup(backend)
    shard_tensor = [platform.from_numpy(numpy.array([1, 2, 3]))] * 3
    global_tensor = platform.from_numpy(numpy.array([1, 1, 1, 2, 2, 2, 3, 3, 3]))
    gather_chunk = CombinationFunc.gather(shard_tensor, dim=0, chunk=3)

    assert platform.allclose(global_tensor, gather_chunk)

tests/test_combination/test_gather.py
alibaba/easydist
# Copyright (c) 2023, Alibaba Group;
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================

import numpy
from easydist import platform
from easydist.metashard.combination import aligned_prefix, shape_aligned_otherdim


def test_aligned_prefix():
    t1 = platform.from_numpy(numpy.array([1, 2, 3, 4]))
    t2 = platform.from_numpy(numpy.array([1, 2, 3, 4]))
    assert 4 == aligned_prefix(t1, t2, dim_idx=0)

    t1 = platform.from_numpy(numpy.array([1, 2, 3, 4]))
    t2 = platform.from_numpy(numpy.array([2, 2, 3, 4]))
    assert 0 == aligned_prefix(t1, t2, dim_idx=0)

    t1 = platform.from_numpy(numpy.array([[1, 2, 3, 4], [1, 2, 3, 4]]))
    t2 = platform.from_numpy(numpy.array([[1, 2, 3, 4], [1, 2, 3, 5]]))
    assert 1 == aligned_prefix(t1, t2, dim_idx=0)
    assert 3 == aligned_prefix(t1, t2, dim_idx=1)


def test_aligned_otherdim():
    shape_1 = (10, 11, 12)
    shape_2 = (10, 13, 12)
    assert shape_aligned_otherdim(shape_1, shape_2, 1) == True
    assert shape_aligned_otherdim(shape_1, shape_2, 2) == False

    shape_1 = (10, 11, 12)
    shape_2 = (10, 13, 13)
    assert shape_aligned_otherdim(shape_1, shape_2, 1) == False
    assert shape_aligned_otherdim(shape_1, shape_2, 2) == False

    shape_1 = (10, 11, 12)
    shape_2 = (10, 11, 12, 13)
    assert shape_aligned_otherdim(shape_1, shape_2, 2) == False
    assert shape_aligned_otherdim(shape_1, shape_2, 3) == False

tests/test_combination/test_help_func.py
alibaba/easydist
# Copyright (c) 2023, Alibaba Group;
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================

import numpy
import pytest

from easydist import platform
from easydist.metashard.combination import CombinationFunc
from easydist.utils.testing import ALL_PLATFORM
from easydist import easydist_setup


@pytest.mark.parametrize("backend", ALL_PLATFORM)
def test_identity(backend):
    easydist_setup(backend)
    shard_tensor = [platform.from_numpy(numpy.array([1, 2, 3]))] * 4
    global_tensor = platform.from_numpy(numpy.array([1, 2, 3]))
    combination_tensor = CombinationFunc.identity(shard_tensor)

    assert platform.allclose(global_tensor, combination_tensor)


@pytest.mark.parametrize("backend", ALL_PLATFORM)
def test_identity_2(backend):
    easydist_setup(backend)
    shard_tensor = [platform.from_numpy(numpy.array([1, 2, 3]))] * 4
    global_tensor = platform.from_numpy(numpy.array([1, 2, 4]))
    combination_tensor = CombinationFunc.identity(shard_tensor)

    assert not platform.allclose(global_tensor, combination_tensor)


@pytest.mark.parametrize("backend", ALL_PLATFORM)
def test_identity_3(backend):
    easydist_setup(backend)
    shard_tensor = [platform.from_numpy(numpy.array([1, 2, 3]))] * 3 + [
        platform.from_numpy(numpy.array([1, 2, 4]))
    ]
    combination_tensor = CombinationFunc.identity(shard_tensor)

    assert combination_tensor is None

tests/test_combination/test_identity.py
alibaba/easydist
# Copyright (c) 2023, Alibaba Group;
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================

import numpy
import pytest

from easydist import platform
from easydist.metashard.combination import CombinationFunc, ReduceOp
from easydist.utils.testing import ALL_PLATFORM
from easydist import easydist_setup


@pytest.mark.parametrize("backend", ALL_PLATFORM)
def test_reduce(backend):
    easydist_setup(backend)
    shard_tensor = [platform.from_numpy(numpy.array([i, i, i])) for i in range(4)]
    max_tensor = platform.from_numpy(numpy.array([3, 3, 3]))
    combination_max = CombinationFunc.reduce(shard_tensor, ops=ReduceOp.MAX)

    assert platform.allclose(max_tensor, combination_max)

    min_tensor = platform.from_numpy(numpy.array([0, 0, 0]))
    combination_min = CombinationFunc.reduce(shard_tensor, ops=ReduceOp.MIN)

    assert platform.allclose(min_tensor, combination_min)

    sum_tensor = platform.from_numpy(numpy.array([6, 6, 6]))
    combination_sum = CombinationFunc.reduce(shard_tensor, ops=ReduceOp.SUM)

    assert platform.allclose(sum_tensor, combination_sum)

tests/test_combination/test_reduce.py
alibaba/easydist
# Copyright (c) 2023, Alibaba Group;
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================

import numpy
import pytest
import functools

from easydist import platform
import easydist.config as mdconfig
from easydist.metashard.combination import CombinationFunc, ReduceOp, try_combination_single
from easydist.utils.testing import ALL_PLATFORM, assert_partial_func_equal
from easydist import easydist_setup


@pytest.mark.parametrize("backend", ALL_PLATFORM)
def test_reduce(backend):
    easydist_setup(backend)
    shard_tensor = [platform.from_numpy(numpy.random.uniform(size=(3, 4))) for _ in range(4)]

    for op_type in [ReduceOp.MAX, ReduceOp.MIN, ReduceOp.SUM]:
        comb_func = functools.partial(CombinationFunc.reduce, ops=op_type)
        global_tensor = comb_func(shard_tensor)

        return_func = try_combination_single(shard_tensor, global_tensor)

        assert_partial_func_equal(comb_func, return_func)


@pytest.mark.parametrize("backend", ALL_PLATFORM)
def test_gather(backend):
    easydist_setup(backend)
    shard_tensor = [platform.from_numpy(numpy.random.uniform(size=(3, 4))) for _ in range(4)]

    for dim_ in [0, 1]:
        comb_func = functools.partial(CombinationFunc.gather, dim=dim_)
        global_tensor = comb_func(shard_tensor)

        return_func = try_combination_single(shard_tensor, global_tensor)

        assert_partial_func_equal(comb_func, return_func)


@pytest.mark.parametrize("backend", ALL_PLATFORM)
def test_gather_halo(backend):
    easydist_setup(backend)
    mdconfig.extend_space = True

    shard_tensor = [platform.from_numpy(numpy.random.uniform(size=(3, 4))) for _ in range(3)]

    for dim_, halo_ in zip([0, 1], [1, 2]):
        comb_func = functools.partial(CombinationFunc.gather, dim=dim_, halowidth=halo_)
        global_tensor = comb_func(shard_tensor)

        return_func = try_combination_single(shard_tensor, global_tensor)

        assert_partial_func_equal(comb_func, return_func)


@pytest.mark.parametrize("backend", ALL_PLATFORM)
def test_gather_chunk(backend):
    easydist_setup(backend)
    mdconfig.extend_space = True

    shard_tensor = [platform.from_numpy(numpy.random.uniform(size=(3, 4))) for _ in range(3)]

    for dim_, chunk_ in zip([0, 1], [3, 2]):
        comb_func = functools.partial(CombinationFunc.gather, dim=dim_, chunk=chunk_)
        global_tensor = comb_func(shard_tensor)

        return_func = try_combination_single(shard_tensor, global_tensor)

        assert_partial_func_equal(comb_func, return_func)

tests/test_combination/test_try_combination_single.py
alibaba/easydist

import logging

import jax
import jax.numpy as jnp

from easydist import easydist_setup, mdconfig
from easydist.jax.api import easydist_compile


@easydist_compile(compile_only=True)
def foo_func(x, y):
    tanh = jnp.tanh(x)
    return jnp.exp(tanh) @ y + tanh


def main():
    mdconfig.log_level = logging.INFO
    easydist_setup(backend="jax", device="cuda")

    key = jax.random.PRNGKey(0)
    key, subkey = jax.random.split(key)
    randn_x = jax.random.normal(key, (10, 10))
    randn_y = jax.random.normal(subkey, (10, 10))

    foo_func(randn_x, randn_y)


if __name__ == '__main__':
    main()

tests/test_strategy/jax/simple_function1.py
alibaba/easydist
import os
import difflib

import torch
from torch._subclasses.fake_tensor import FakeTensorMode

from easydist import easydist_setup
from easydist.torch.experimental.api import easydist_compile
from easydist.torch import set_device_mesh
from easydist.utils.testing import TorchMockDeviceMesh


class Foo(torch.nn.Module):

    def __init__(self):
        super().__init__()
        self.norm = torch.nn.LayerNorm(5)
        self.linear = torch.nn.Linear(5, 5)

    def forward(self, x):
        x = self.norm(x)
        x = self.linear(x)
        return x.relu()


@easydist_compile(tracing_mode="fake", cuda_graph=False, compile_only=True)
def train_step(input, model, opt):
    out = model(input).mean()
    out.backward()
    opt.step()
    opt.zero_grad(True)

    return out


def train_example():
    fake_mode = FakeTensorMode()

    torch.ones(1).cuda()
    with torch.device('cuda'), fake_mode:
        model = Foo()
        randn_input = torch.randn(16, 5)

        torch.distributed.broadcast(randn_input, src=0)

        opt = torch.optim.Adam(model.parameters(), lr=0.001, foreach=True, capturable=True)

    # trace train step func
    mock_mesh = TorchMockDeviceMesh(1, 2, debug_only=True)
    set_device_mesh(mock_mesh)

    train_step(randn_input, model, opt)

def main():
    # setting up easydist and torch.distributed
    easydist_setup(backend="torch", device="cuda", allow_tf32=False)

    torch.distributed.init_process_group(backend="nccl")
    local_rank = int(os.environ["LOCAL_RANK"])
    torch.cuda.set_device(local_rank)

    train_example()

if __name__ == "__main__":
    main()


tests/test_strategy/torch/simple_model1.py
alibaba/easydist
# Copyright (c) 2023, Alibaba Group;
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================

import pytest
import torch
import functorch
from functorch.compile import aot_function
import rich

from easydist.utils.testing import TorchMockDeviceMesh
from easydist.torch import EDTorchShardingAnn, set_device_mesh
from easydist.torch.passes import fix_addmm_bias, eliminate_detach
from easydist import easydist_setup


def fn_1(x, y):
    return torch.concat([x, y], dim=1)


def fn_2(x, y):
    return torch.mm(torch.exp(torch.tanh(x)), y)


@functorch.compile.make_boxed_compiler
def compiler_fn(fx_module: torch.fx.GraphModule, inps):
    fx_module = fix_addmm_bias(fx_module)
    fx_module = eliminate_detach(fx_module)
    fx_module.recompile()
    print(fx_module.graph)

    sharding_interpreter = EDTorchShardingAnn(fx_module)
    sharding_info, fwd_shape_info = sharding_interpreter.run(*inps)
    rich.print("sharding_info:\n", sharding_info)
    rich.print("fwd_shape_info:\n", fwd_shape_info)

    return fx_module


@pytest.mark.parametrize("fn", [fn_1, fn_2])
def test_simple_case(fn):
    easydist_setup("torch")

    mock_mesh = TorchMockDeviceMesh(1, 2, debug_only=True)
    set_device_mesh(mock_mesh)

    x = torch.randn(10, 10, requires_grad=True)
    y = torch.randn(10, 10, requires_grad=True)
    aot_print_fn = aot_function(fn, fw_compiler=compiler_fn, bw_compiler=compiler_fn)
    res = aot_print_fn(x, y)

    grad_res = torch.ones_like(res)
    res.backward(grad_res)


if __name__ == '__main__':
    test_simple_case()

tests/test_torch/test_simple.py
alibaba/easydist
# Copyright (c) 2023, Alibaba Group;
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================

import functools

import torch

from easydist.metashard.combination import CombinationFunc
from easydist.metashard import ShardAnnotation, ShardDim, MetaOp
from easydist.utils.testing import assert_partial_func_equal
from easydist import easydist_setup


def test_metaop_preset():
    easydist_setup("torch")
    input_args = (torch.rand((3, 4, 768)), 3, 2), {}
    meta_op = MetaOp(torch.ops.aten.chunk, input_args)
    preset_anno = ShardAnnotation([[ShardDim(0), ShardDim(0), ShardDim(1, chunk=3)]])
    comb_func = meta_op.sharding_discovery_with_preset(preset_anno)

    right_answer = [functools.partial(CombinationFunc.gather, dim=2)] * 3

    assert comb_func != None
    assert len(comb_func) == len(right_answer)

    for func1, func2 in zip(comb_func, right_answer):
        assert_partial_func_equal(func1, func2)

tests/test_unfiyshard/test_unifyop.py
alibaba/easydist
# Copyright (c) 2023, Alibaba Group;
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================

import functools

from easydist.metashard import ShardAnnotation, ShardDim
from easydist.metashard.combination import CombinationFunc
from easydist.metashard.view_propagation import view_propagation_preset
from easydist.utils.testing import assert_partial_func_equal


def test_view_propagation_preset():
    preset_anno = ShardAnnotation([[ShardDim(1, chunk=5), ShardDim(0)]])
    comb_func = view_propagation_preset([10, 8], [5, 2, 8], preset_anno)
    answer = functools.partial(CombinationFunc.gather, dim=1)
    assert_partial_func_equal(comb_func, answer)

    preset_anno = ShardAnnotation([[ShardDim(0), ShardDim(1, chunk=2)]])
    comb_func = view_propagation_preset([10, 8], [10, 2, 2, 2], preset_anno)
    answer = functools.partial(CombinationFunc.gather, dim=2)
    assert_partial_func_equal(comb_func, answer)

    preset_anno = ShardAnnotation([[ShardDim(0), ShardDim(1, chunk=4)]])
    comb_func = view_propagation_preset([10, 8], [10, 2, 2, 2], preset_anno)
    answer = functools.partial(CombinationFunc.gather, dim=3)
    assert_partial_func_equal(comb_func, answer)

    preset_anno = ShardAnnotation([[ShardDim(1, chunk=3), ShardDim(0)]])
    comb_func = view_propagation_preset([10, 8], [5, 2, 8], preset_anno)
    assert comb_func is None

    preset_anno = ShardAnnotation([[ShardDim(0), ShardDim(1, chunk=2)]])
    comb_func = view_propagation_preset([10, 8], [5, 2, 8], preset_anno)
    assert comb_func is None


test_view_propagation_preset()

tests/test_unfiyshard/test_view_propagation.py
alibaba/easydist