Add --dataloader_multiprocessing_context to work around Python 3.14 incompatibility - #9941
Open
sliedes wants to merge 1 commit into
Open
Add --dataloader_multiprocessing_context to work around Python 3.14 incompatibility#9941sliedes wants to merge 1 commit into
sliedes wants to merge 1 commit into
Conversation
Collaborator
|
Thanks for the contribution! Based on your code, we've attempted a fix for both the template pickling and CUDA initialization issues — could you check whether this works for your use case: https://github.com/modelscope/ms-swift/pull/9979/changes |
Author
|
Thanks! It seems that has already been merged in main? I tested current main (087aa1c) and still see this. Running: Output: |
Collaborator
Thanks for testing! Let me check |
Collaborator
|
Seems I missed some attributes in |
Author
Yes, this seems to fix it for me! Side note, feel free to ignore: I did need to update my single-line JPEG XL support plugin because with import multiprocessing
import pillow_jxl
multiprocessing.set_forkserver_preload(["pillow_jxl"]) |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
PR type
PR information
Background
In Python 3.14, the default multiprocessing method on Linux changed from
forkto the much saferforkserver. This, however, exposes a latent issue in swift (see #4586): The data loader has a reference to the entire model, which Python tries to pickle and send to the workers.It seems someone has previously taken care to not serialize the entire model:
ms-swift/swift/pipelines/train/sft.py
Line 309 in 302b074
This protects one serialization path, but after restoring
template.model,__post_process_datasets(), wraps the dataset usingtemplate.encodeas a bound method which makes the model reachable again:ms-swift/swift/pipelines/train/sft.py
Line 137 in 302b074
This was bad already before Python 3.14; CUDA and
forkdo not go well together. However, it has worked in practice because the data loader doesn't touch anything CUDA.In practice, when using Python 3.14, I have used
--dataloader_num_workers 0to work around this problem; however, that's obviously not ideal.What would a proper fix look like
I think a proper fix would ensure that the model is not reachable from the DataLoader so that it can be properly pickled. Perhaps distinguishing a worker template from a model-bound training template; conceptually
The DataLoader's dataset and collator get the worker-safe one. Model-dependent
_post_encodehappens in the main training process with the real one.What this PR does
I am hesitant to do bigger refactorings of a code base I do not understand well enough, so this PR provides what I consider a minimal reasonable workaround for the current state of things: It allows specifying an optional parameter
--dataloader_multiprocessing_context, with valid optionsfork,forkserver,spawn, defaulting toNone, which I believe is the reasonable safe choice as this really only papers over a bug. Now, with--dataloader_multiprocessing_context forkmy training runs work on Python 3.14.There's a wart: Transformers do not, in released versions, expose a parameter to do this; hence, the PR overrides
_get_dataloaderand changes the loader's held context. Inmain, transformers has adopted this parameter:https://github.com/huggingface/transformers/blob/f9b76f2dfc44fdc6109468925bf0e1856fd34278/src/transformers/training_args.py#L1318
Reflections
I'm sending this hope it's useful. I would be more than happy if someone actually fixed the underlying issue of the data loader having a reference to the model (and thus being unpickleable) instead. I hope this at least demonstrates what the issue is even if deemed too hacky to merge.