현대 대안: jax.sharding과 Mesh

2024년 이후 JAX 팀은 jax.sharding과 Mesh를 새로운 분산 API로 권장해. pmap보다 깨끗하고, 모델 병렬 처리까지 자연스럽게 확장해.

핵심 아이디어는 배열 자체에 "어떻게 분산되어 있는지"라는 sharding 정보를 부여. 그러면 jit이 sharding을 알아서 해석하고 집합 통신을 자동 삽입해.

import jax
import jax.numpy as jnp
from jax.sharding import Mesh, PartitionSpec as P, NamedSharding
import numpy as np

# ============ Mesh 만들기 ============
# 8 devices 를 (data=4, model=2) 2D mesh 로
devices = np.array(jax.devices()[:8]).reshape(4, 2)
mesh = Mesh(devices, axis_names=("data", "model"))

# ============ Sharded array 만들기 ============
sharding = NamedSharding(mesh, P("data", "model"))

# 큰 array 를 mesh 에 분산
x = jnp.arange(64*32).reshape(64, 32)
x_sharded = jax.device_put(x, sharding)

# 시각화
jax.debug.visualize_array_sharding(x_sharded)
# +-----+-----+
# | d0  | d4  |   ← data=0, model=0,1
# |-----+-----|
# | d1  | d5  |
# |-----+-----|
# | d2  | d6  |
# |-----+-----|
# | d3  | d7  |
# +-----+-----+

jit이 자동으로 분산 처리

@jax.jit
def matmul(a, b):
    return a @ b

# input 을 sharded 하면 — output 도 sharded, collective 자동 삽입
a = jax.device_put(jnp.ones((1024, 512)), NamedSharding(mesh, P("data", None)))
b = jax.device_put(jnp.ones((512, 256)), NamedSharding(mesh, P(None, "model")))

c = matmul(a, b)  # output 도 (data, model) 분산

같은 함수, sharding만 바뀌면 데이터 병렬 처리, 텐서 병렬 처리, FSDP 모두 가능해.

학습 루프, sharding으로

def train_step(params, x, y, lr):
    def loss_fn(p):
        return jnp.mean((x @ p - y) ** 2)
    loss, grads = jax.value_and_grad(loss_fn)(params)
    new_params = jax.tree.map(lambda p, g: p - lr * g, params, grads)
    return new_params, loss

train_step = jax.jit(train_step)

# data parallel: params 는 replicated, batch 는 sharded
data_sharding = NamedSharding(mesh, P("data", None))
replicated = NamedSharding(mesh, P())   # 모든 device 가 같은 값

params_repl = jax.device_put(params, replicated)
x_sharded = jax.device_put(batch_x, data_sharding)
y_sharded = jax.device_put(batch_y, NamedSharding(mesh, P("data")))

# jit 이 자동으로 collective 삽입
new_params, loss = train_step(params_repl, x_sharded, y_sharded, 0.01)

pmap 대비 장점:

jit 호환: 평범한 jit으로 분산할 수 있어. pmap의 별도 wrapping이 없어.
다양한 sharding: 데이터 병렬 처리, 모델 병렬 처리, FSDP, 파이프라인 등을 같은 API로 다뤄.
집합 통신 자동: pmean / psum / all_gather 등을 명시하지 않아도 돼. XLA가 자동으로 처리해.
합성 가능: vmap / grad와 자유롭게 합성해.

🆕 새 코드는 sharding으로

2024년 이후 JAX 코드는 pmap보다 sharding API를 써. pmap은 사용 중단 상태는 아니지만 유지보수 모드야. DeepMind와 Google의 새 모델들(Gemini 등)은 모두 sharding 기반이야. 이 퀘스트의 다음 레슨도 sharding 패턴으로 다뤄.

pmap 코드를 sharding으로 옮길 때는 다음 순서로 진행해.

jax.pmap(f, axis_name="data") → jax.jit(f)
매개변수 replicate → NamedSharding(mesh, P())
batch shard → NamedSharding(mesh, P("data"))
jax.lax.pmean(grads) → 자동 (jit이 처리)

Code

import jax
import jax.numpy as jnp
from jax.sharding import Mesh, NamedSharding, PartitionSpec as P
# Since JAX 0.7+, you can also use: jax.P as a shortcut for PartitionSpec

# Step 1: Create a mesh — a logical arrangement of devices
devices = jax.devices()  # All available devices
mesh = jax.make_mesh((len(devices),), ('data',))

# Step 2: Define sharding with PartitionSpec
# P('data') means "shard along the 'data' axis of the mesh"
data_sharding = NamedSharding(mesh, P('data'))
replicated = NamedSharding(mesh, P())  # No sharding = replicated

# Step 3: Place data on devices
X = jax.random.normal(jax.random.PRNGKey(0), (128, 4))
X_sharded = jax.device_put(X, data_sharding)

# Check the sharding
jax.debug.visualize_array_sharding(X_sharded)

import jax
import jax.numpy as jnp

# Create a 2D mesh: (data_parallel, model_parallel)
# For example, 8 devices arranged as 4x2
mesh = jax.make_mesh((4, 2), ('dp', 'mp'))

# Data is sharded along the 'dp' axis
data_sharding = jax.NamedSharding(mesh, jax.P('dp', None))

# Model weights are sharded along the 'mp' axis
weight_sharding = jax.NamedSharding(mesh, jax.P(None, 'mp'))

# With jit + sharding constraints, JAX handles communication automatically
@jax.jit
def forward(params, x):
    return x @ params

# JAX inserts the necessary all-gathers and reduce-scatters

import jax
import jax.numpy as jnp

# Since JAX 0.7, jax.P is an alias for jax.sharding.PartitionSpec
mesh = jax.make_mesh((4, 2), ('dp', 'mp'))

# These are equivalent:
sharding1 = jax.NamedSharding(mesh, jax.P('dp', None))
sharding2 = jax.NamedSharding(mesh, jax.sharding.PartitionSpec('dp', None))

# jax.P is more concise and is the recommended style

🆕 새 코드는 sharding으로

Code

External links

Exercise

Progress

댓글 0