Py314t uses 2× memory

Python Help
,
1

I’m migrating BEVFormer (an auto-driving model training framework) from Py3.14 to 3.14t. Functionally it runs, but memory usage (RAM, not GPU memory) on 3.14t is about twice that of 3.14. I’ve run many training epochs and tried different GPUs; the issue reproduces reliably.

I traced it to BRC (biased ref count): when a non-owner thread decrefs and ob_ref_shared comes to 0, it has to have the owner thread read ob_ref_local and confirm the object can be collected. If the owner is detached, collection is delayed and the object sits in the BRC queue.

In torch.DataLoader this case is common: owner threads stay detached for a long time, tied to how multiprocessing.Queue is implemented. A repro is here.

import gc, os, re, sys, time
import torch, torch.multiprocessing as mp
HDR = re.compile(r"^[0-9a-f]+-[0-9a-f]+ ")
NB = 64 * 1024 * 1024

# count mem usage
def mb():
    tot, path = 0.0, ""
    for l in open("/proc/self/smaps"):
        if HDR.match(l):
            f = l.split(); path = f[5] if len(f) > 5 else ""
        elif l.startswith("Size:") and path.startswith("/dev/shm/torch"):
            tot += int(l.split()[1]) / 1024
    return tot

def consumer(q):
    while True:
        x = q.get()
        if x is None: 
            break
        del x

if __name__ == "__main__":
    q = mp.Queue()
    p = mp.Process(target=consumer, args=(q,))
    p.start()
    gc.collect(); gc.disable()
    t = torch.empty(NB)
    q.put(t)
    del t
    time.sleep(1)
    [None for _ in range(3 * 10**7)] # spin
    print(f"  GIL={sys._is_gil_enabled()}  {mb():5.0f} MB", flush=True)
    q.put(None)
    p.join()

How should this be fixed? The most direct approach is to change torch.DataLoader or multiprocessing.Queue so the owner does not stay detached. I tried that; it needs a lot of code changes and didn’t work well.

Another way is to change BRC in CPython so a non-owner can move the owner from detached to suspended (similar to STW), read ob_ref_local, decide whether the object can be collected, then put the owner back from suspended to detached. That sounds quite unsafe, but it worked: on BEVFormer again, 3.14t memory matched 3.14.

Is there a better way to handle this? I can post my BRC patch; it was written on 3.14, not current main.

2

I forgot to mention GC. Yes, GC can collect obj in BRC queue. In this case, though, GC almost never runs: the memory Torch allocates is mostly large, self-managed objects, not lots of small ones, so GC is hard to trigger.

3

What you’re experiencing is definitely one of the expected downsides of biased reference counting in free-threading, but 2x memory is quite concerning indeed.

Can you walk through how the queue is used in the system? Though, I suspect changing from a multiprocessing queue to a multithreading queue is not on the table?

I don’t know what suspended means here. Can you elaborate further on the changes you made?

4

Hi. detached and suspended are thread state. The state transition diagram can be found in Include/internal/pycore_pystate.h

//            (bound thread)        (stop-the-world thread)
// [attached]       <->       [detached]       <->       [suspended]
//   |                                                        ^
//   +---------------------------->---------------------------+
//                          (bound thread)
5

Hoestly I havn’t take a closer look at how the queue is used in pytorch. A colleague of mine is more familiar with pytorch, let me discuss this with him

6

Oh I see, I was thinking about object states sorry.

I think that’s quite an intrusive change, though. I don’t think it would be good to have that behavior in CPython.

One thing you could try is something I was tinkering with when implementing a multithreaded queue, but I believe I didn’t push it anywhere, sorry. The idea is that if the thread that puts an object in a queue is also the object owner (which is fairly likely), then you can disown the object before incrementing the shared reference count. (I believe there should be something like this somewhere in a header file in main; sorry I don’t have access to my laptop right now.) This puts the object in the merged state, and thus it can be deallocated without sending it through the BRC queue.

I should have time to look into this next week. If you can confirm whether this fixes the memory usage problem that would be great.

7

Great idea! Given that queues are usually used to pass objects across threads, this makes a lot of sense. I believe the function is _Py_ExplicitMergeRefcount. I’ll try this approach tomorrow.

8

It doesn’t help you until PyTorch adds Python 3.15/3.15t support, but Python 3.15 might do better on this benchmark.

@nas did some work for Python 3.15 that should make the GC more efficient at reducing memory use in some situations. I’m not sure off-hand if this case is the same as what’s covered by Neil’s work.

9

I don’t think QSBR impacts this. The objects in question should be in the default state, not the maybe-weakrefs one, IIUC.

10

I had Claude build me a reproducer that doesn’t need pytorch. Note that calling gc.collect() explicitly will “drain” the BRC queue. Perhaps that would be a mitigation to this issue?

Note that this change was not backported to 3.14. The memory being held by the data loader is perhaps not allocated by mimalloc. So, I would guess that change doesn’t have much effect, whether reverted or not.

There is the idea of having other threads drain the BRC queue, on behalf of a detached owner. That seems tricky to get right though. The owner may be reattaching concurrently, ob_ref_local is written non-atomically by the owner, and the merging thread needs a safe way to know the owner will stay detached for the duration. The existing stop-the-world machinery gives a correct but expensive version of this, and something like triggering it only once the queue exceeds a size threshold might be a reasonable middle ground. Draining the queue on reattach would be cheap and worth doing regardless, though it does not help a thread that stays blocked.

11

I read the reproducer, and I believe this is the issue I meet in bevformer.

I also looked at the revert patch you mentioned, and I agree that it is unlikely to help much. (but I think it looks worthy to backport to 3.14)

something like triggering it only once the queue exceeds a size threshold might be a reasonable middle ground

I considered a similar approach too, i.e. still relying on GC to collect. Length of the BRC queue alone may not be enough, though; the sizes of the queued objects also matter. At least in my case, the objects sitting in the BRC queue are large.

12

@dpdani Hi, I tried your approach — disown the object on enqueue — but found that in my case it probably won’t help. The queue here is a multiprocessing.Queue, which is inter-process, not inter-thread. (Yes, torch.DataLoader is multi-process rather than multi-thread.) So objects have to be serialized when enqueue. For torch.tensors, obviously nobody want to serialize such a large chunk of memory, so Torch uses a complex trick to do 0-copy sharing. And I believe that’s where the problem is introduced.

To summarize: for queue, disown the object enqueued or disown the object produced by serialization won’t solve the issue I ran into. That said, I still think the idea (disown the object on enqueue) is a great idea, or at least it’s worth offering as an option that users can choose to enable.

13

Hi @weixlu, I’ve tried running your script and I can reproduce the ~2x memory use, compared to GIL-enabled 3.14. I can also reproduce Neil’s script.

I still think the culprit is in multiprocessing.Queue, which uses an internal thread to send data through the pipe (see multiprocessing.queues.Queue._feed) and a collections.deque as internal buffer. If I change deque_append_impl in _collectionsmodule.c to use the _Py_ExplicitMergeRefcount trick instead of Py_NewRef I can see the RSS memory usage go down, though not back to what it was under the GIL-enabled build. Whereas, the /dev/shm/torch memory is back down to 0.

I don’t think we should apply this change to CPython because deque is definitely used in single-threaded contexts which should not see a slowdown, but can you confirm it would be helpful to your case?

14

Wow, it works on the reproducer! :star_struck: Thanks a lot for digging into this.

Sure, I will start the model training and see if it helps in real-world case

15

@dpdani It turns out be effective, but “disown” is more complex than we expected. In the reproducer, each enqueue is a single tensor. In the real workload, each enqueue is something like (task_id, {"img": [tensor], "other_data": "other_data"}), so disowning only the top-level object is not enough — you also have to walk its references recursively to find the tensors.

That walk looks a lot like GC, and I used gc.get_referents() for it. Reference cycles make it messier; my current approach is using a set of already-seen object ids (which makes it sound even more like GC). That extra work increases enqueue overhead, so I’m not sure this would be a general-purpose queue.

But at least, in my case it works. I trained on 8 x A100 GPUs with 4 dataloader processes per GPU (32 processes total). Baseline (py3.14) used 10GB; py3.14t used 20GB; with the recursive-disown queue it came down to 14GB.

16

@weixlu I have a fix for this at GitHub - kumaraditya303/cpython at backport-5778062-3.14 · GitHub for 3.14 branch. Can you please try that? I changed the brc queue to force merge refcounts for detached threads if there are more than 10 objects in the queue.

17

Ah yes, thanks! Sure, let me give it a try

18

Hi, thanks for your fix. I verified this can resolve the memory-usage issue. But for my case, _Py_BRC_MERGE_THRESHOLD needs to be set very low. Py3.14 (as baseline) uses 10GB memory, Py3.14t uses 20GB. With your fix, it drops to 10GB when _Py_BRC_MERGE_THRESHOLD=1, and to 17GB when the threshold is 2.

19

Thanks for verifying it. I was expecting this because we don’t have a good way to figure out object size so relying on object count is not perfect measure of memory pressure here. Still, I think it is good that this patch works for you with low threshold.

20

@weixlu This is now fixed in main branch with gh-157838: Merge biased refcounts on behalf of detached threads by kumaraditya303 · Pull Request #157839 · python/cpython · GitHub and backported to 3.14 and 3.15 branches.