Open southfreebird opened 2 months ago
southfreebird
commented
2 months ago I've found the documentation about unmapped inputs and outputs, and it seems like it addresses the problem I'm fighting against. However, I can't figure out if there is a way to make it work with Pallas.
mattjj
commented
1 month ago Sorry for the slow response!
One workaround is not to use unmapped inputs or outputs in shard_map: instead, do any broadcasts or sums you need outside the shard_map, and rely on XLA to generate the right code automatically (i.e. turning explicit broadcasting or jnp.tile calls on inputs into no-ops, and turning jnp.sums to reduction collectives).
Another workaround is not to use autodiff at all, and to use a custom_vjp around the shard_map.
I find both of those workarounds pretty unsatisfactory!
A better solution would be to, as you describe, make the pallas_call work with shard_map's check_rep=True. Logically that means telling shard_map's internals how your particular pallas_call affects cross-device replication; then, once shard_map can track cross-device replication through it, we can set check_rep=True and thus avoid defensive allreduces (as described in the doc you linked). (We'll also avoid the error "No replication rule for pallas_call", which I mention here so that this issue is more discoverable.) Unfortunately we don't yet have public APIs for registering that information.
Until we have a public API for this, here's a sketch of how to do it using JAX-internal APIs, which may change without warning:
MeshAxisName = str
from jax._src.pallas.pallas_call import pallas_call_p
from jax.experimental.shard_map import register_norewrite, register_check
@register_check(pallas_call_p)
def pallas_check(
mesh: Mesh, *in_rep: set[MeshAxisName], name: str, **params
) -> list[set[MeshAxisName]]:
del params # ignore other pallas_call_p params
if name == '_copy_kernel_pallas':
inputs_rep, = in_rep # single input in this case
return inputs_rep # no change in replication, single output
else:
raise NotImplementedError(name)
register_norewrite(pallas_call_p)I inlined it into your file and also set check_rep=True on the pallas_call, then got this output:
HloModule jit_grad_func, is_scheduled=true, entry_computation_layout={()->bf16[1,128]{1,0}}, allow_spmd_sharding_propagation_to_output={true}, num_partitions=8
ENTRY %main.36_spmd () -> bf16[1,128] {
%constant.9 = bf16[] constant(1)
ROOT %broadcast.4 = bf16[1,128]{1,0} broadcast(bf16[] %constant.9), dimensions={}, metadata={op_name="jit(grad_func)/jit(main)/convert_element_type[new_dtype=bfloat16 weak_type=False sharding=None]" source_file="/usr/local/google/home/mattjj/packages/jax/28155.py" source_line=82}
}What do you think?
mattjj
commented
1 month ago Actually, a better rule for pallas calls that don't communicate might look something like this:
from typing import Sequence
MeshAxisName = str
from jax._src.pallas.pallas_call import pallas_call_p
from jax.experimental.shard_map import (register_rewrite, register_check,
_standard_rewrite_rule)
@register_check(pallas_call_p)
def pallas_check(
mesh: Mesh, *in_rep: set[MeshAxisName],
name: str, out_shapes: Sequence[jax.ShapeDtypeStruct], **_,
) -> list[set[MeshAxisName]]:
num_outputs = len(out_shapes)
if name == '_copy_kernel_pallas':
in_rep_ = [r for r in in_rep if r is not None]
if in_rep_ and not in_rep_[:-1] == in_rep_[1:]:
raise Exception("rewrite rule failed")
return [in_rep_[0]] * num_outputs
else:
raise NotImplementedError(name)
@register_rewrite(pallas_call_p)
def pallas_rewrite(mesh, in_rep, *args, **params):
name = params['name']
if name == '_copy_kernel_pallas':
return _standard_rewrite_rule(pallas_call_p, mesh, in_rep, *args, **params)
else:
raise NotImplementedError(name)Oops, this issue isn't closed; I closed the wrong issue with that PR.
Description
Hello team, I have found unexpected behavior when using
shard_map,custom_vjp, andpallas(tritonworks the same) simultaneously during auto grad. Please take a look at the code snippet:Output:
As you can see, in the backward pass, when my output is replicated by 2 GPUs (output sharding does not contain the
hiddenaxis, so it's not parallelized by themodel), I encounter an unexpected AllReduce. As I understand it, this is needed to sync data by themodelMesh axis between 2 GPUs. However, in my particular case, I don't need this synchronization and want to get rid of it, but I can't find a good solution for this. It doesn't reproduce if we don't have a custom_call. Could you help me with that?Thank you!
System info (python version, jaxlib version, accelerator, etc.)
jax: 0.4.28 jaxlib: 0.4.28 numpy: 1.24.3 python: 3.10.12 (main, Nov 20 2023, 15:14:05) [GCC 11.4.0] jax.devices (8 total, 8 local): [cuda(id=0) cuda(id=1) ... cuda(id=6) cuda(id=7)] process_count: 1 platform: uname_result(system='Linux', release='5.4.0-155-generic', version='#172-Ubuntu SMP Fri Jul 7 16:10:02 UTC 2023', machine='x86_64')