Comments (6)
Do you have a minimal code example that can reproduce?
Just profiling jax.block_until_ready(jax.numpy.zeros((9500000,3), dtype=jnp.float32))
in isolation takes <500us on my end on a V100.
from jax.
Oh, XLA JIT compilation should be used to reproduce. Thank you! I can repro a compute time latency of 131 microseconds with
@jax.jit
def run_zeros() -> jax.Array:
"""Run the main logic to benchmark."""
with jax.named_scope("benchmark_broadcast"):
return jnp.zeros((9500000, 3), dtype=jnp.float32)
values = run_zeros()
jax.block_until_ready(values)
print(values)
The resulting HLO indicates the slow broadcast:
@@f32[9500000,3]{1,0} fusion(), kind=kLoop, calls=
{
tmp_0 = f32[] constant(0)
ROOT tmp_1 = f32[9500000,3]{1,0} broadcast(f32[] tmp_0), dimensions={}
}
from jax.
Might it be possible that the CUDA implementation for broadcast here only uses 1 SM? And is thus not parallelized?
from jax.
Ah, it's 130 microseconds, not milliseconds?
I don't think it's physically possible to do better?
A V100 has 900GB/s memory bandwidth, per the specs. 9500000 * 3 * 4 bytes / 900e9 bytes/sec = 127us
.
(I'm a bit surprised here, actually, I thought NVIDIA's published bandwidth numbers are bidirectional, which would mean twice as long, but your measurement seems to indicate otherwise.)
So... unless I'm mistaken, you're almost getting roofline memory bandwidth and you should not expect to do better?
from jax.
Ohhh... Apologies. It is 131 microseconds! I have corrected the original post.
from jax.
Thank you for the analysis! I suppose allocating a new memory block (versus reading from existing memory) also encounters this memory bandwidth. That would present a bottleneck indeed that means we can't do better.
from jax.
Related Issues (20)
- `random.key_impl`wrong return type annotation HOT 1
- Bug in single-precision linear system solving HOT 4
- error: failed to legalize operation 'mhlo.erf' with jax-metal
- jnp.vectorize causes ```FloatingPointError: invalid value (nan)``` when there is no nan returned. HOT 4
- AttributeError: module 'jax' has no attribute 'linear_util' HOT 4
- [Pallas – Triton] add `tl.debug_barrier` equivalent HOT 1
- `NNFunctionsTest.testDotProductAttention` failing on A100/H100 with `use_vmap=True, `impl='cudnn'` HOT 2
- _isolve tacitly checks for symmetric linear operator wrongly HOT 4
- Improve CuSolver errors diagnostics HOT 8
- CuSolver: Switch to 64 bit api to allow for eigh on matrices > than 26732x26732 HOT 6
- `jnp.arange` does not permit dynamic shaped arrays HOT 2
- Profile shows TPU operations before CPU instructions
- ValueError: Apply must be called with an RNG as the third argument, the required signature is: `apply(params, state, rng, *a, **k)`. The object was of type <class 'jax._src.interpreters.partial_eval.DynamicJaxprTracer'>: Traced<ShapedArray(uint32[2])>with<DynamicJaxprTrace(level=1/0) HOT 5
- Cannot export with certain symbolic args constraints HOT 1
- Provide a dummy ```jax.key(...)``` that allows splitting, but not sampling HOT 3
- TPU operations don't appear if profile is too long
- Multi-process GPU jobs fail on Slurm HOT 6
- Error exporting np.where when unrelated dimension variable is 0 HOT 9
- Operator `jax.lax.gather` consumes a lot of memory on CPU HOT 6
- Force a physical axis layout rearrangement HOT 3
Recommend Projects
-
React
A declarative, efficient, and flexible JavaScript library for building user interfaces.
-
Vue.js
🖖 Vue.js is a progressive, incrementally-adoptable JavaScript framework for building UI on the web.
-
Typescript
TypeScript is a superset of JavaScript that compiles to clean JavaScript output.
-
TensorFlow
An Open Source Machine Learning Framework for Everyone
-
Django
The Web framework for perfectionists with deadlines.
-
Laravel
A PHP framework for web artisans
-
D3
Bring data to life with SVG, Canvas and HTML. 📊📈🎉
-
Recommend Topics
-
javascript
JavaScript (JS) is a lightweight interpreted programming language with first-class functions.
-
web
Some thing interesting about web. New door for the world.
-
server
A server is a program made to process requests and deliver data to clients.
-
Machine learning
Machine learning is a way of modeling and interpreting data that allows a piece of software to respond intelligently.
-
Visualization
Some thing interesting about visualization, use data art
-
Game
Some thing interesting about game, make everyone happy.
Recommend Org
-
Facebook
We are working to build community through open source technology. NB: members must have two-factor auth.
-
Microsoft
Open source projects and samples from Microsoft.
-
Google
Google ❤️ Open Source for everyone.
-
Alibaba
Alibaba Open Source for everyone
-
D3
Data-Driven Documents codes.
-
Tencent
China tencent open source team.
from jax.