tinygrad / tinygrad/tinygrad

sharding empty uses double peak mem and residual mem

Open
#13,513 1 comment 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

Dominant language
Python
Stars
33.6k
Forks
4.3k
Avg merge
3h 11m
Merged PRs (30d)
491

Description

from tinygrad import Tensor, Device
ds = tuple(f"{Device.DEFAULT}:{i}" for i in range(8))
t = Tensor.empty(10**9).shard(ds, 0).realize()
scheduled 16 kernels in 6.19 ms (221 uops in cache)
*** AMD        1 E_1953125_16_4                                 arg  2 mem   4.50 GB tm    478.56us/     0.48ms (      0 GFLOPS 2090|2090   GB/s) 
*** AMD        2 E_1953125_16_4n1                               arg  2 mem   5.00 GB tm    483.28us/     0.96ms (      0 GFLOPS 2069|2069   GB/s) 
*** AMD        3 E_1953125_16_4n2                               arg  2 mem   5.50 GB tm    485.12us/     1.45ms (      0 GFLOPS 2061|2061   GB/s) 
*** AMD        4 E_1953125_16_4n3                               arg  2 mem   6.00 GB tm    486.72us/     1.93ms (      0 GFLOPS 2055|2055   GB/s) 
*** AMD        5 E_1953125_16_4n4                               arg  2 mem   6.50 GB tm    486.88us/     2.42ms (      0 GFLOPS 2054|2054   GB/s) 
*** AMD        6 E_1953125_16_4n5                               arg  2 mem   7.00 GB tm    489.72us/     2.91ms (      0 GFLOPS 2042|2042   GB/s) 
*** AMD        7 E_1953125_16_4n6                               arg  2 mem   7.50 GB tm    487.40us/     3.40ms (      0 GFLOPS 2052|2052   GB/s) 
*** AMD        8 E_1953125_16_4n7                               arg  2 mem   8.00 GB tm    492.20us/     3.89ms (      0 GFLOPS 2032|2032   GB/s) 
*** AMD        9 E_1953125_16_4n8                               arg  2 mem   4.50 GB tm    491.52us/     4.38ms (      0 GFLOPS 2035|2035   GB/s) 
*** AMD:1     10 xfer  500.00M,   AMD:1 <- AMD                  arg  2 mem   4.50 GB tm   8730.90us/    13.11ms (      0 GFLOPS   57|57     GB/s) 
*** AMD:2     11 xfer  500.00M,   AMD:2 <- AMD                  arg  2 mem   4.50 GB tm   8731.92us/    21.84ms (      0 GFLOPS   57|57     GB/s) 
*** AMD:3     12 xfer  500.00M,   AMD:3 <- AMD                  arg  2 mem   4.50 GB tm   8710.56us/    30.55ms (      0 GFLOPS   57|57     GB/s) 
*** AMD:4     13 xfer  500.00M,   AMD:4 <- AMD                  arg  2 mem   4.50 GB tm   8702.11us/    39.26ms (      0 GFLOPS   57|57     GB/s) 
*** AMD:5     14 xfer  500.00M,   AMD:5 <- AMD                  arg  2 mem   4.50 GB tm   8687.91us/    47.94ms (      0 GFLOPS   58|58     GB/s) 
*** AMD:6     15 xfer  500.00M,   AMD:6 <- AMD                  arg  2 mem   4.50 GB tm   8707.93us/    56.65ms (      0 GFLOPS   57|57     GB/s) 
*** AMD:7     16 xfer  500.00M,   AMD:7 <- AMD                  arg  2 mem   4.50 GB tm   8703.59us/    65.36ms (      0 GFLOPS   57|57     GB/s) 

because it creates empty on device 0 first then copies each shard, the peak mem is 8GB and ends with 4.5GB. this should creates empty on each device and uses 4GB total

Contributor guide

No contributing guide indexed for this repository

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Reproduce the issue with the Tensor.empty(...).shard(ds, 0).realize() example and inspect the Tensor.empty, shard, and realize paths, including device transfers. Trace why allocation occurs on the default device before shards are copied. Done means each device allocates its shard directly, peak memory avoids the extra full tensor, and the residual total is about 4GB.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
distributed-systems, performance
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Stale
Clarity
Mostly clear
Newbie friendliness
42/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.