sharding.py 13.3 KB
Newer Older
1
# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2
3
#
# See LICENSE for license information.
4
5
6
7
8
9
10
"""Sharding utilities for Transformer Engine in JAX.

This module provides utilities for managing tensor sharding in distributed training,
including support for various parallelism strategies like data parallelism (DP),
tensor parallelism (TP), pipeline parallelism (PP), and full-sharded data
parallelism (FSDP). It includes functions for sharding constraints, mesh management,
and collective operations.
11
"""
12
import os
13
14
15
from contextlib import contextmanager
from dataclasses import dataclass
from enum import Enum
16
from typing import Callable, Optional
17
import warnings
18
19
import jax
import jax.numpy as jnp
20
21
from jax.interpreters import pxla
from jax.sharding import PartitionSpec, get_abstract_mesh
22
import numpy as np
23
24
25

_PXLA_THREAD_RESOURCES = pxla.thread_resources

26
# Axis Names
27
28
29
BATCH_AXES = "nvte_batch"
SEQLEN_AXES = "nvte_seqlen"
SEQLEN_TP_AXES = "nvte_seqlen_tp"
30
SEQLEN_CP_AXES = "nvte_seqlen_cp"
31
32
33
34
35
36
37
38
HEAD_AXES = "nvte_head"
HIDDEN_AXES = "nvte_hidden"
HIDDEN_TP_AXES = "nvte_hidden_tp"
JOINED_AXES = "nvte_joined"
W_NO_SHARD_AXES = "nvte_w_no_shard"
W_FSDP_AXES = "nvte_w_fsdp"
W_TP_AXES = "nvte_w_tp"
W_JOINED_AXES = "nvte_w_joined"
39

40

41
def _get_mesh_info(resource: str, mesh: jax.sharding.Mesh):
42
    assert resource in mesh.axis_names, f"{resource} is not in the axis_names of Mesh {mesh}."
43
44
45
    return mesh.shape[resource], resource


46
47
48
49
50
51
52
53
def get_sharding_map_logic_axis_to_mesh_axis():
    """
    Generate a dict to map logical axes to mesh axes.
    """
    gsr = global_mesh_resource()

    IS_FSDP_OUTER = bool(int(os.environ.get("NVTE_OUTER_BATCH_FSDP_DIM", False)))

54
55
56
57
58
    batch_resources = (
        [gsr.fsdp_resource, gsr.dp_resource]
        if IS_FSDP_OUTER
        else [gsr.dp_resource, gsr.fsdp_resource]
    )
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75

    batch_dim_rule = []
    for resource in batch_resources:
        if resource is not None and resource not in batch_dim_rule:
            batch_dim_rule.append(resource)

    if len(batch_dim_rule) <= 0:
        batch_dim_rule = None
    elif len(batch_dim_rule) == 1:
        batch_dim_rule = batch_dim_rule[0]
    else:
        batch_dim_rule = tuple(batch_dim_rule)

    te_logical_axis_to_mesh_axis = {
        BATCH_AXES: batch_dim_rule,
        SEQLEN_AXES: None,
        SEQLEN_TP_AXES: gsr.tp_resource,
76
        SEQLEN_CP_AXES: gsr.cp_resource,
77
78
79
80
81
82
83
84
85
86
87
88
        HEAD_AXES: gsr.tp_resource,
        HIDDEN_AXES: None,
        HIDDEN_TP_AXES: gsr.tp_resource,
        JOINED_AXES: None,
        W_NO_SHARD_AXES: None,
        W_FSDP_AXES: gsr.fsdp_resource,
        W_TP_AXES: gsr.tp_resource,
        W_JOINED_AXES: None,
    }
    return te_logical_axis_to_mesh_axis


89
def generate_pspec(logical_axis_names, with_flax_rules=False, padded=False):
90
91
92
    """
    Convert logical axes to PartitionSpec
    """
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
    rules = None
    if with_flax_rules:
        try:
            import flax

            rules = dict(flax.linen.get_logical_axis_rules())
        except ImportError:
            pass

    if rules is None:
        warnings.warn(
            "Transformer Engine logical axes, such as BATCH_AXES, SEQLEN_AXES, etc. are deprecated"
            " and removed in a future version. Please use Flax logical axes with the"
            " `flax.linen.logical_axis_rules()` context and optionally use"
            " `transformer_engine.jax.flax.extend_logical_axis_rules()` to extend Flax axis rules"
            " with Transformer Engine logical axes.",
            DeprecationWarning,
        )
        rules = get_sharding_map_logic_axis_to_mesh_axis()
112
113
114
115
116
    # mesh_axis_names = [rules[name] for name in logical_axis_names]
    mesh_axis_names = []
    for name in logical_axis_names:
        axis_name = rules[name] if name in rules else None
        mesh_axis_names.append(axis_name)
117
    pspec = jax.sharding.PartitionSpec(*mesh_axis_names)
118
119
    if padded:
        pspec = get_padded_spec(pspec, len(mesh_axis_names))
120
121
122
123
124
    return pspec


def with_sharding_constraint(x: jnp.array, pspec: PartitionSpec):
    """
125
126
127
128
    A wrapper function to jax.lax.with_sharding_constraint
        1. Does nothing if mesh is empty.
        2. If all mesh axes are manual axes, replaces pspec with all Nones.
        3. Otherwise, strips only the manual axes.
129
130
131
132
133
134
135
    """
    if pspec is None:
        return x

    mesh = _PXLA_THREAD_RESOURCES.env.physical_mesh
    if mesh.empty:
        return x
136
137
138
139
140
141
142
143

    # We want to exclude the axes that already used by shard_map and shard_map
    # only sets those in the abstract_mesh, not the physical one
    manual_axis_names = get_abstract_mesh().manual_axes
    cleaned_axis_names = tuple(name if name not in manual_axis_names else None for name in pspec)

    cleaned_pspec = PartitionSpec(*cleaned_axis_names)
    return jax.lax.with_sharding_constraint(x, cleaned_pspec)
144
145


146
147
148
def with_sharding_constraint_by_logical_axes(
    x: jnp.array, logical_axis_names: Optional[tuple | list]
):
149
    """
150
151
152
    A wrapper function to flax.linen.with_logical_constraint.

    DEPRECATED USE CASE: If no Flax logical axis rules are available, this function falls back to jax.lax.with_sharding_constraint using a hardcoded logical axis rule table from TE rules, such as BATCH_AXES. This functionality will be removed in the future.
153
154
155
156
157
158
159
160
161
162
163

    If logical_axis_names = None, this means no sharding constraint is applied.

    If logical_axis_names = (None, None, ...), this means a sharding constraint is applied and the tensor is replicated across all devices.

    Args:
        x: Input tensor to apply sharding constraint
        logical_axis_names: Logical axis names to apply sharding constraint
    Returns:
        Tensor with sharding constraint applied, or the original tensor if no logical axes are provided.

164
    """
165
    if not logical_axis_names:
166
167
        return x

168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
    try:
        # Check if Flax logical axis rules are available, if so use them
        import flax

        flax_rules = flax.linen.get_logical_axis_rules()
        if len(flax_rules) > 0:
            return flax.linen.with_logical_constraint(
                x, logical_axis_names, fallback=flax.linen.spmd.RulesFallback.NO_CONSTRAINT
            )
    except ImportError:
        pass

    warnings.warn(
        "TransformerEngine logical axes, such as BATCH_AXES, SEQLEN_AXES, etc. are deprecated and"
        " will be removed in a future version. Please use Flax logical axes with a"
        " flax.linen.logical_axis_rules context and optionally use"
        " transformer_engine.jax.flax.extend_logical_axis_rules to add BATCH_AXES, etc. to your"
        " rules.",
        DeprecationWarning,
    )

    # If no logical axis rules are available from Flax, fallback to TE's hardcoded logical axis rule table
190
191
192
193
194
    assert len(x.shape) == len(logical_axis_names)
    pspec = generate_pspec(logical_axis_names)
    return with_sharding_constraint(x, pspec)


195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
def get_all_mesh_axes():
    """
    Get all name of mesh axes
    """
    mesh = _PXLA_THREAD_RESOURCES.env.physical_mesh
    return mesh.axis_names


def get_padded_spec(spec, ndim):
    """
    Get padded spec for partitioning from arguments' information
    """
    if spec is None:
        return (None,) * ndim
    assert len(spec) <= ndim
    return spec + (None,) * (ndim - len(spec))


213
214
215
def lax_paral_op(
    x: jnp.array, ops: Callable, mesh_resource: str, mesh: jax.sharding.Mesh, **kwargs
):
216
217
218
219
    """
    A wrapper function to invoke lax.p* operations, like psum.
    """
    if mesh_resource is not None:
220
        _, resource = _get_mesh_info(mesh_resource, mesh)
221
        return ops(x, resource, **kwargs)
222
223
224
225
226
227
228
229
230
231
    return x


def num_of_devices():
    """
    Get total number of detected devices
    """
    return len(jax.devices())


232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
def get_mesh_axis_size(axis, mesh=None):
    """
    Get the axis size of the given mesh.
    If the mesh is None, it would be replaced
    by the global mesh.
    """
    if mesh is None:
        mesh = _PXLA_THREAD_RESOURCES.env.physical_mesh

    if axis is None:
        return 1

    assert axis in mesh.shape, f"{axis} is not a axis of the given mesh {mesh.shape}"
    return mesh.shape[axis]


def get_mesh_axis_rank(axis: str, mesh=None):
    """
    Gets the local axis rank of the `axis` of the array.
    If the mesh is None the rank is 0.
    """
    if mesh is None:
        return 0
    _, axis_name = _get_mesh_info(axis, mesh)
    return jax.lax.axis_index(axis_name)


259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
def get_mesh_axis_rank_host(axis, mesh) -> int:
    """
    Same as get_mesh_axis_rank(), but return a host value instead of a
    traced device value.
    """
    if axis not in mesh.axis_names:
        raise ValueError(f"Axis {axis} not found in mesh axis names: {mesh.axis_names}")

    axis_index = mesh.axis_names.index(axis)

    # Convert mesh.devices (ndarray of Device objects) to flat list
    devices = mesh.devices
    local_device = jax.devices()[jax.process_index()]  # Pick one device on this host

    # Find index of local_device in mesh.devices
    coords = np.argwhere(devices == local_device)
    if coords.size == 0:
        raise ValueError(f"Local device {local_device} not found in mesh.devices.")
    coords = tuple(coords[0])  # Coordinates in the mesh array

    # Get the mesh rank along the specified axis
    rank = coords[axis_index]
    return int(rank)


284
@dataclass
285
class MeshResource:
286
287
288
289
290
291
292
293
294
295
296
    """A data container for managing mesh resources in distributed training.

    This class defines the mapping between logical axes and physical mesh axes
    for different types of parallelism in distributed training.

    Attributes:
        dp_resource: Axis name for data parallelism (batch sharding), default is None
        tp_resource: Axis name for tensor parallelism (hidden dimension sharding), default is None
        fsdp_resource: Axis name for full-sharded data parallelism, default is None
        pp_resource: Axis name for pipeline parallelism (layer sharding), default is None
        cp_resource: Axis name for context parallelism (sequence sharding), default is None
297
    """
298

299
300
    dp_resource: str = None
    tp_resource: str = None
301
    fsdp_resource: str = None
302
    pp_resource: str = None
303
    cp_resource: str = None
304
305


306
_GLOBAL_MESH_RESOURCE = MeshResource()
307
308
309


@contextmanager
310
def global_shard_guard(resource: MeshResource):
311
312
313
314
315
316
317
    """Context manager for setting global sharding configuration.

    This context manager allows temporarily setting the global mesh resource
    configuration for sharding operations.

    Args:
        resource: MeshResource instance defining the sharding configuration
318
    """
319
    global _GLOBAL_MESH_RESOURCE
320
    old_resources = _GLOBAL_MESH_RESOURCE
321
    try:
322
        _GLOBAL_MESH_RESOURCE = resource
323
324
        yield
    finally:
325
        _GLOBAL_MESH_RESOURCE = old_resources
326

327

328
def global_mesh_resource() -> MeshResource:
329
330
331
332
    """Get the current global mesh resource configuration.

    Returns:
        The current MeshResource instance
333
334
    """
    return _GLOBAL_MESH_RESOURCE
335

336

337
def all_reduce_sum_along_dp_fsdp(x: jnp.array, mesh: jax.sharding.Mesh):
338
339
340
341
342
343
344
345
    """Perform all-reduce sum operation along data parallelism and FSDP axes.

    Args:
        x: Input tensor to reduce
        mesh: JAX mesh for distributed computation

    Returns:
        Reduced tensor
346
    """
347
348
    x = lax_paral_op(x, jax.lax.psum, global_mesh_resource().dp_resource, mesh)
    return lax_paral_op(x, jax.lax.psum, global_mesh_resource().fsdp_resource, mesh)
349
350


351
def all_reduce_max_along_all_axes_except_PP(x: jnp.array, mesh: jax.sharding.Mesh):
352
353
354
355
356
357
358
359
    """Perform all-reduce max operation along all axes except pipeline parallelism.

    Args:
        x: Input tensor to reduce
        mesh: JAX mesh for distributed computation

    Returns:
        Reduced tensor
360
361
362
363
    """
    all_axes = get_all_mesh_axes()
    for axis in all_axes:
        if axis != global_mesh_resource().pp_resource:
364
            x = lax_paral_op(x, jax.lax.pmax, axis, mesh)
365
366
367
368
369
370
371
    return x


# Deprecating Items ---------------------------------------------------------------
ShardingResource = MeshResource

global_shard_resource = global_mesh_resource
372
373
374


class MajorShardingType(Enum):
375
376
377
378
379
380
381
382
383
384
    """Enumeration of major sharding types for distributed training.

    This enum defines the basic sharding patterns available for distributed
    training. Note that this class is deprecated and will be removed in the future.

    Values:
        SINGLE: Single process training
        DP: Data parallel training
        TP: Standard tensor parallel training
        DPTP: Data and standard tensor parallel training
385
    """
386

387
388
389
390
391
392
393
    SINGLE = 0
    DP = 1
    TP = 2
    DPTP = 3


class ShardingType(Enum):
394
395
396
397
398
399
400
401
402
403
404
405
406
    """Enumeration of detailed sharding types for distributed training.

    This enum defines specific sharding patterns for distributed training,
    including combinations of data parallelism and different tensor parallelism
    strategies. Note that this class is deprecated and will be removed in the future.

    Values:
        SINGLE: No sharding
        DP: Sharding along data parallelism
        TP_COL: Sharding along column-split tensor parallelism
        TP_ROW: Sharding along row-split tensor parallelism
        DP_TP_COL: Sharding along data and column-split tensor parallelism
        DP_TP_ROW: Sharding along data and row-split tensor parallelism
407
    """
408

409
410
411
412
413
414
    SINGLE = (MajorShardingType.SINGLE, "single")
    DP = (MajorShardingType.DP, "dp")
    TP_COL = (MajorShardingType.TP, "tp_col")
    TP_ROW = (MajorShardingType.TP, "tp_row")
    DP_TP_COL = (MajorShardingType.DPTP, "dp_tp_col")
    DP_TP_ROW = (MajorShardingType.DPTP, "dp_tp_row")