Support configurable multiprocessing start method (forkserver/spawn) to avoid fork deadlocks

[Thaurun]

Problem

The dataloader can deadlock with the default fork start method in environments where the
parent process already has other threads holding locks (OpenMP, MKL, CUDA, HuggingFace tokenizers,
logging handlers, etc.). fork() only copies the calling thread, so locks held elsewhere stay
permanently locked in the worker → deadlock.

Solution

Make the multiprocessing start method configurable (fork / forkserver / spawn). forkserver/
spawn start a fresh interpreter and avoid inherited locks. Default stays fork — no behavior change.

Changes

Configurable start method

• loader.py: get_loader / get_savable_loader accept multiprocessing_context: str = "fork".
• savable_loader.py: loaders create queues + torch DataLoader context from
  torch.multiprocessing.get_context(mp_ctx); config() now includes multiprocessing_context.
• worker.py: WorkerConfig gains __getstate__/__setstate__ that drops the non-picklable
  data_parallel_group (ProcessGroup) on serialization.

Make dataset callables picklable (required because forkserver/spawn serialize the dataset to workers)

• default_generic_webdataset.py: lambdas → functools.partial of _part_in_set, _apply_field_map, _wrap_sample.
• joined_webdataset.py: joiner lambda → functools.partial(_join_samples, ...).
• crude.py: default lambdas → module-level _always_true_part / _identity_sample.
• mix_batch_dataset.py: default batch_mix_fn → _identity_batch_mix.
• iter_map_dataset.py: default len_map_fn → _identity_len.

Supporting

• base_dataset.py: _function_config renders functools.partial as readable functools.partial(<module>.<qualname>).
• tests/test_forkserver.py (new): save/restore + fork-order equivalence under forkserver.
• tests/test_metadataset.py: updated reference config() for the new multiprocessing_context key and sample_loader repr.

Usage

loader = get_savable_loader(dataset, multiprocessing_context="forkserver")  # or "spawn"

[VVsssssk]

Does capturing subflavors=subflavors or {} here pose a behavior regression risk?

If dataset.yaml doesn't have subflavors, but the caller/Metadataset subsequently executes dataset.subflavors.update(subflavors) at src/megatron/energon/dataset_config.py:138, the subflavors in the samples won't see the update.

[Thaurun]

Good catch — yes, that was a regression. When subflavors is None, the two subflavors or {} expressions created distinct empty dicts, so later dataset.subflavors.update(...) no longer affected subflavors on samples. Fixed by resolving the dict once and sharing that same object between self.subflavors and the picklable _wrap_sample partial.

[VVsssssk]

Only data_parallel_group was removed — should _worker_debug_file also be removed? It may not be picklable either.

When worker_log_level >= 1 and worker_debug_path is set, the loader initialization opens the debug file in the main process; the dataset serialization may then fail when forkserver/spawn serializes it.

[Thaurun]

Agreed. With worker_log_level >= 1 and worker_debug_path set, SavableDataLoader.init already opens _worker_debug_file on the shared WorkerConfig before workers are spawned, so forkserver/spawn would try to pickle an open file handle. Fixed by also dropping _worker_debug_file (and _worker_debug_file_worker_id) in getstate; workers reopen the file lazily in worker_log.

[Thaurun]

Addressed all review comments — fixed the shared subflavors dict, dropped _worker_debug_file from pickle state, and made FileStoreCachePool picklable for forkserver/spawn.