diff --git a/.github/workflows/acceptancetest_mg7.yml b/.github/workflows/acceptancetest_mg7.yml index bc334e43b..f293c43ef 100644 --- a/.github/workflows/acceptancetest_mg7.yml +++ b/.github/workflows/acceptancetest_mg7.yml @@ -201,6 +201,116 @@ jobs: export PATH="$HOME/.cache/HEPtools/bin:$PATH" ./tests/test_manager.py test_group_subprocess_mg7 -pA -t0 -l INFO + acceptancetest_mg7_vegas_reproducibility: + needs: build_madspace + # mg7 seeded-generation reproducibility (p p > t t~, plain VEGAS-optimized + # run): the same run_card seed must give a byte-identical LHE file across + # independent runs, a different seed must not. Self-skips if the madspace + # + LHAPDF(NNPDF23) stack is unavailable. + runs-on: ubuntu-24.04 + if: github.event_name == 'push' || github.event.pull_request.head.repo.fork == true + steps: + - uses: actions/checkout@v4 + - uses: ./.github/actions/checkout_mg5 + - uses: ./.github/actions/install_madspace + - uses: ./.github/actions/restore-pip-cache + - uses: ./.github/actions/restore_heptools + - name: test one of the test test_vegas_reproducibility_mg7 + run: | + cd $GITHUB_WORKSPACE + export PATH="$HOME/.cache/HEPtools/bin:$PATH" + ./tests/test_manager.py test_vegas_reproducibility_mg7 -pA -t0 -l INFO + + acceptancetest_mg7_madnis_reproducibility: + needs: build_madspace + # mg7 seeded-generation reproducibility for a full survey + madnis + # training + generate run (p p > t t~, no gridpack export): the same + # run_card seed must give a byte-identical LHE file across independent + # runs, including the training phase itself, a different seed must not. + # Companion to acceptancetest_mg7_gridpack_reproducibility, which only + # covers a frozen gridpack's own regeneration. Self-skips if the madspace + # + LHAPDF(NNPDF23) stack is unavailable. + runs-on: ubuntu-24.04 + if: github.event_name == 'push' || github.event.pull_request.head.repo.fork == true + steps: + - uses: actions/checkout@v4 + - uses: ./.github/actions/checkout_mg5 + - uses: ./.github/actions/install_madspace + - uses: ./.github/actions/restore-pip-cache + - uses: ./.github/actions/restore_heptools + - name: test one of the test test_madnis_reproducibility_mg7 + run: | + cd $GITHUB_WORKSPACE + export PATH="$HOME/.cache/HEPtools/bin:$PATH" + ./tests/test_manager.py test_madnis_reproducibility_mg7 -pA -t0 -l INFO + + acceptancetest_mg7_madnis_reproducible_mode: + needs: build_madspace + # mg7 madnis.reproducible = True (p p > t t~): unlike + # acceptancetest_mg7_madnis_reproducibility, this requires madnis training + # itself to be byte-identical for the same seed independent of the cpu + # thread pool size, both with online-only training and with buffered + # (off-policy replay) training enabled. Self-skips if the madspace + + # LHAPDF(NNPDF23) stack is unavailable. + runs-on: ubuntu-24.04 + if: github.event_name == 'push' || github.event.pull_request.head.repo.fork == true + steps: + - uses: actions/checkout@v4 + - uses: ./.github/actions/checkout_mg5 + - uses: ./.github/actions/install_madspace + - uses: ./.github/actions/restore-pip-cache + - uses: ./.github/actions/restore_heptools + - name: test one of the test test_madnis_reproducible_mode_mg7 + run: | + cd $GITHUB_WORKSPACE + export PATH="$HOME/.cache/HEPtools/bin:$PATH" + ./tests/test_manager.py test_madnis_reproducible_mode_mg7 -pA -t0 -l INFO + + acceptancetest_mg7_gridpack_reproducibility: + needs: build_madspace + # mg7 seeded-generation reproducibility for a gridpack trained with + # madnis (p p > t t~): the gridpack's own bin/generate_events --seed must + # give a byte-identical LHE file across independent runs with the same + # seed, and a different result for a different seed. Training itself is + # not required to be deterministic. Self-skips if the madspace + + # LHAPDF(NNPDF23) stack is unavailable. + runs-on: ubuntu-24.04 + if: github.event_name == 'push' || github.event.pull_request.head.repo.fork == true + steps: + - uses: actions/checkout@v4 + - uses: ./.github/actions/checkout_mg5 + - uses: ./.github/actions/install_madspace + - uses: ./.github/actions/restore-pip-cache + - uses: ./.github/actions/restore_heptools + - name: test one of the test test_gridpack_reproducibility_mg7 + run: | + cd $GITHUB_WORKSPACE + export PATH="$HOME/.cache/HEPtools/bin:$PATH" + ./tests/test_manager.py test_gridpack_reproducibility_mg7 -pA -t0 -l INFO + + acceptancetest_mg7_gridpack_reproducibility_vegas: + needs: build_madspace + # mg7 seeded-generation reproducibility for a plain VEGAS-optimized + # gridpack (p p > t t~, no madnis training): the gridpack's own + # bin/generate_events --seed must give a byte-identical LHE file across + # independent runs with the same seed, and a different result for a + # different seed. Companion to acceptancetest_mg7_gridpack_reproducibility, + # which covers the madnis-trained case. Self-skips if the madspace + + # LHAPDF(NNPDF23) stack is unavailable. + runs-on: ubuntu-24.04 + if: github.event_name == 'push' || github.event.pull_request.head.repo.fork == true + steps: + - uses: actions/checkout@v4 + - uses: ./.github/actions/checkout_mg5 + - uses: ./.github/actions/install_madspace + - uses: ./.github/actions/restore-pip-cache + - uses: ./.github/actions/restore_heptools + - name: test one of the test test_gridpack_reproducibility_vegas_mg7 + run: | + cd $GITHUB_WORKSPACE + export PATH="$HOME/.cache/HEPtools/bin:$PATH" + ./tests/test_manager.py test_gridpack_reproducibility_vegas_mg7 -pA -t0 -l INFO + acceptancetest_mg7_merged_flavor_uq: needs: build_madspace # mg7 cross-section for the merged-flavor u q > u q (q = u d), pinned to the diff --git a/madgraph/iolibs/template_files/mg7/gridpack.py b/madgraph/iolibs/template_files/mg7/gridpack.py index 62626f327..ccc71999b 100644 --- a/madgraph/iolibs/template_files/mg7/gridpack.py +++ b/madgraph/iolibs/template_files/mg7/gridpack.py @@ -31,7 +31,6 @@ import tomllib import argparse - def resolve_verbosity(verbosity: str) -> str: """Resolve the run_card "auto" verbosity to "pretty"/"log" depending on whether stdout is attached to a terminal; other values pass through @@ -69,6 +68,12 @@ def main() -> None: # parse command line arguments parser = argparse.ArgumentParser() parser.add_argument("--run_name", type=str, default=run_args["run_name"]) + parser.add_argument( + "--seed", type=int, default=run_args.get("seed", 0), + help="0 draws a fresh random seed; a positive integer makes the run " + "reproducible at any --cpu_thread_pool_size (for --output_format lhe " + "also set combine_thread_pool_size = 1)" + ) parser.add_argument("--device", type=str, nargs="*") parser.add_argument( "--cpu_thread_pool_size", type=int, default=run_args["cpu_thread_pool_size"] @@ -173,6 +178,7 @@ def main() -> None: channels=channel_generators, status_file=os.path.join(run_path, "info.json"), config=config, + seed=args.seed, ) # run generation diff --git a/madgraph/iolibs/template_files/mg7/madevent.py b/madgraph/iolibs/template_files/mg7/madevent.py index 2ab4b2f4c..97d1c8ba0 100644 --- a/madgraph/iolibs/template_files/mg7/madevent.py +++ b/madgraph/iolibs/template_files/mg7/madevent.py @@ -203,7 +203,7 @@ def init_context(self) -> None: self.device_types = [] self.devices = [] self.pool_sizes = [] - for i, device_name in enumerate(device_names): + for device_name in device_names: if ":" in device_name: device_type, device_index_str = device_name.split(":") device_index = int(device_index_str) @@ -417,6 +417,7 @@ def build_event_generator(self, phasespaces: list[PhaseSpace]) -> ms.EventGenera channels=channel_generators, status_file=os.path.join(self.run_path, "info.json"), config=self.event_generator_config, + seed=self.run_card["run"]["seed"], ) unused_globals = ( set(self.contexts[0].global_names()) - event_generator.used_globals() @@ -427,36 +428,42 @@ def build_event_generator(self, phasespaces: list[PhaseSpace]) -> ms.EventGenera return event_generator def survey_phasespaces( - self, phasespaces: list[PhaseSpace | None] + self, phasespaces: list[PhaseSpace | None], survey_pass: int = 0 ) -> ms.EventGenerator | None: ps_filtered = [ps for ps in phasespaces if ps is not None] if len(ps_filtered) == 0: return None event_generator = self.build_event_generator(ps_filtered) - event_generator.survey() + event_generator.survey(survey_pass) return event_generator def survey(self) -> None: + # survey_pass distinguishes the survey() calls below: "both" mode can + # re-survey a channel carried over unchanged from the multichannel pass + # into the final (simplified) pass, and both passes schedule jobs on the + # same underlying ChannelEventGenerator. The explicit pass index keeps each + # pass's job seeds independent of the other passes' job counts, rather than + # depending on call history. phasespace_mode = self.run_card["phasespace"]["mode"] if phasespace_mode == "multichannel": self.phasespaces = [ subproc.build_multichannel_phasespace() for subproc in self.subprocesses ] - self.event_generator = self.survey_phasespaces(self.phasespaces) + self.event_generator = self.survey_phasespaces(self.phasespaces, 0) elif phasespace_mode == "flat": self.phasespaces = [ subproc.build_flat_phasespace() for subproc in self.subprocesses ] - self.event_generator = self.survey_phasespaces(self.phasespaces) + self.event_generator = self.survey_phasespaces(self.phasespaces, 0) elif phasespace_mode == "both": kept_count = self.run_card["phasespace"]["simplified_channel_count"] phasespaces_multi = [ subproc.build_multichannel_phasespace() for subproc in self.subprocesses ] - evgen_multi = self.survey_phasespaces(phasespaces_multi) + evgen_multi = self.survey_phasespaces(phasespaces_multi, 0) phasespaces_flat = [ subproc.build_flat_phasespace() @@ -464,7 +471,7 @@ def survey(self) -> None: None for subproc in self.subprocesses ] - #evgen_flat = self.survey_phasespaces(phasespaces_flat, "flat") + #evgen_flat = self.survey_phasespaces(phasespaces_flat, 1) channel_status = evgen_multi.channel_status() cross_sections = [] @@ -485,7 +492,7 @@ def survey(self) -> None: self.subprocesses, phasespaces_multi, phasespaces_flat, cross_sections ) ] - self.event_generator = self.survey_phasespaces(self.phasespaces) + self.event_generator = self.survey_phasespaces(self.phasespaces, 2) else: raise ValueError("Unknown phasespace mode") @@ -525,6 +532,7 @@ def train_madnis(self) -> None: config.buffer_unweighting_quantile = madnis_args["buffer_unweighting_quantile"] config.fixed_cwnet_fraction = madnis_args["fixed_cwnet_fraction"] config.softclip_threshold = madnis_args["softclip_threshold"] + config.reproducible = madnis_args["reproducible"] madnis_phasespaces = [] integrands = [] cwnets = [] @@ -550,6 +558,12 @@ def train_madnis(self) -> None: config=config, integrands=integrands, cwnets=cwnets, + # Reuses the run_card's own seed (also used by build_event_generator()); + # MultiMadnisTraining derives an independent, non-colliding stream from it + # (see salt::job_kind_madnis_train in random.hpp). Only the single-channel + # CPU sample-generation path is currently seeded -- buffered training and + # GPU multi-channel batches are still non-deterministic. + seed=run_args["seed"], ) madnis_training.train() for phasespace, active_channels in zip( @@ -1292,6 +1306,12 @@ def simplify_phasespace( def build_madnis(self, phasespace: PhaseSpace) -> PhaseSpace: madnis_args = self.process.run_card["madnis"] + # Shared across all networks below: each one's global names are + # already unique (per-channel/per-component prefixes), so + # initialize_globals() derives an independent, non-colliding stream + # per tensor from this one base seed (see global_init_seed in + # random.hpp) -- no extra indices need to be threaded through here. + seed = self.process.run_card["run"]["seed"] channels = [] for channel_id, channel in enumerate(phasespace.channels): prefix = f"subproc{self.subproc_id}.channel{channel_id}" @@ -1309,7 +1329,7 @@ def build_madnis(self, phasespace: PhaseSpace) -> PhaseSpace: subnet_layers=madnis_args["discrete_layers"], subnet_activation=self.activation(madnis_args["discrete_activation"]), ) - discrete_before.initialize_globals(self.process.contexts[0]) + discrete_before.initialize_globals(self.process.contexts[0], seed) cond_dim += perm_count flow_dim = channel.phasespace_mapping.random_dim() @@ -1324,10 +1344,11 @@ def build_madnis(self, phasespace: PhaseSpace) -> PhaseSpace: invert_spline=madnis_args["flow_invert_spline"], ) if channel.adaptive_mapping is None: - flow.initialize_globals(self.process.contexts[0]) + flow.initialize_globals(self.process.contexts[0], seed) else: flow.initialize_from_vegas( - self.process.contexts[0], channel.adaptive_mapping.grid_name() + self.process.contexts[0], channel.adaptive_mapping.grid_name(), + seed ) cond_dim += flow_dim @@ -1342,7 +1363,7 @@ def build_madnis(self, phasespace: PhaseSpace) -> PhaseSpace: subnet_layers=madnis_args["discrete_layers"], subnet_activation=self.activation(madnis_args["discrete_activation"]), ) - discrete_after.initialize_globals(self.process.contexts[0]) + discrete_after.initialize_globals(self.process.contexts[0], seed) channels.append(Channel( phasespace_mapping = channel.phasespace_mapping, @@ -1411,7 +1432,9 @@ def build_cwnet(self, channel_count: int) -> ms.ChannelWeightNetwork: activation=self.activation(madnis_args["cwnet_activation"]), prefix=f"subproc{self.subproc_id}.cwnet", ) - cwnet.initialize_globals(self.process.contexts[0]) + cwnet.initialize_globals( + self.process.contexts[0], self.process.run_card["run"]["seed"] + ) return cwnet def t_channel_mode(self, name: str) -> ms.PhaseSpaceMapping.TChannelMode: diff --git a/madgraph/iolibs/template_files/mg7/run_card.toml b/madgraph/iolibs/template_files/mg7/run_card.toml index 28410c262..35b9bc173 100644 --- a/madgraph/iolibs/template_files/mg7/run_card.toml +++ b/madgraph/iolibs/template_files/mg7/run_card.toml @@ -1,5 +1,11 @@ [run] run_name = %(run.run_name)s +# 0 draws a fresh random seed each run; a positive integer makes the run +# reproducible: the same (seed, thread-pool sizes) reproduces bit-identically and +# distinct seeds are statistically independent. Generation is reproducible at any +# cpu_thread_pool_size; for output_format = lhe also set combine_thread_pool_size = 1 +# (compact_npy and lhe_npy are reproducible at any combine size). +seed = %(run.seed)s # options: cuda, hip, cpp, cppnone, cppsse4, cppavx2, cpp512y, cpp512z, cppauto devices = %(run.devices)s # options: @@ -136,3 +142,4 @@ batch_size_threshold = %(madnis.batch_size_threshold)s channel_grouping_mode = %(madnis.channel_grouping_mode)s # options: none, uniform, learned fixed_cwnet_fraction = %(madnis.fixed_cwnet_fraction)s softclip_threshold = %(madnis.softclip_threshold)s +reproducible = %(madnis.reproducible)s # make training fully reproducible for a given seed, independent of thread count (CPU only, some throughput/memory cost) diff --git a/madgraph/various/banner.py b/madgraph/various/banner.py index 5d5b0c63d..474d6465e 100755 --- a/madgraph/various/banner.py +++ b/madgraph/various/banner.py @@ -6453,6 +6453,12 @@ def default_setup(self): # ----------------------------- [run] -------------------------- self.add_toml_param('run', 'run_name', "run", gridpack=True) + self.add_toml_param('run', 'seed', 0, gridpack=True, + comment="0 draws a fresh random seed each run; any positive integer " + "makes the run reproducible: the same (seed, thread-pool sizes) " + "reproduces bit-identically. Generation is reproducible at any " + "cpu_thread_pool_size; for output_format = lhe also set " + "combine_thread_pool_size = 1") self.add_toml_param('run', 'devices', ["cppnone"], typelist=str, gridpack=True, comment="options: cuda, hip, cpp, cppnone, cppsse4, cppavx2, cpp512y, cpp512z, cppauto") self.add_toml_param('run', 'simd_vector_size', -1, @@ -6588,6 +6594,7 @@ def default_setup(self): allowed=['none', 'uniform', 'learned']) self.add_toml_param('madnis', 'fixed_cwnet_fraction', 0.33) self.add_toml_param('madnis', 'softclip_threshold', 30.0) + self.add_toml_param('madnis', 'reproducible', False) # ----------------- dynamic (free-form) sections --------------- self.dynamic_sections['multiparticles'] = collections.OrderedDict([ diff --git a/madspace/CMakeLists.txt b/madspace/CMakeLists.txt index 342fd11b4..5d87bd1be 100644 --- a/madspace/CMakeLists.txt +++ b/madspace/CMakeLists.txt @@ -136,6 +136,8 @@ if(ENABLE_OPENBLAS) -DBUILD_TESTING=OFF -DBUILD_WITHOUT_CBLAS=ON -DUSE_THREAD=OFF + -DUSE_LOCKING=ON + -DDYNAMIC_ARCH=ON -DCMAKE_POSITION_INDEPENDENT_CODE=1 INSTALL_BYPRODUCTS /lib/${openblas_libname} ) @@ -178,6 +180,7 @@ add_library( src/driver/backend.cpp src/driver/lhe_output.cpp src/driver/madnis_training.cpp + src/driver/random.cpp src/compgraphs/type.cpp src/compgraphs/function.cpp src/compgraphs/instruction.cpp @@ -230,6 +233,7 @@ add_library( include/madspace/driver/lhe_output.hpp include/madspace/driver/logger.hpp include/madspace/driver/madnis_training.hpp + include/madspace/driver/random.hpp include/madspace/compgraphs.hpp include/madspace/compgraphs/type.hpp include/madspace/compgraphs/function.hpp diff --git a/madspace/include/madspace/driver/backend.hpp b/madspace/include/madspace/driver/backend.hpp index 792d5828b..91ff8e935 100644 --- a/madspace/include/madspace/driver/backend.hpp +++ b/madspace/include/madspace/driver/backend.hpp @@ -1,5 +1,8 @@ #pragma once +#include +#include + #include "madspace/compgraphs.hpp" #include "madspace/driver/context.hpp" #include "madspace/driver/tensor.hpp" @@ -9,9 +12,14 @@ namespace madspace { class Runtime { public: virtual ~Runtime() = default; - virtual TensorVec run(const TensorVec& inputs) = 0; + // `seed`, when set, makes this call's random numbers a deterministic function of + // `seed` alone. Backends that can't support this should reject it, not ignore it. + virtual TensorVec + run(const TensorVec& inputs, std::optional seed = std::nullopt) = 0; virtual std::tuple> run_with_grad( - const TensorVec& inputs, const std::vector& input_requires_grad + const TensorVec& inputs, + const std::vector& input_requires_grad, + std::optional seed = std::nullopt ) = 0; virtual std::pair run_backward( const TensorVec& output_grads, @@ -19,6 +27,7 @@ class Runtime { const std::vector& eval_grad, bool return_contiguous_grads = false ) = 0; + friend std::unique_ptr build_runtime(const Function& function, ContextPtr context, bool concurrent); diff --git a/madspace/include/madspace/driver/channel_generator.hpp b/madspace/include/madspace/driver/channel_generator.hpp index 957037d15..48e2c4d71 100644 --- a/madspace/include/madspace/driver/channel_generator.hpp +++ b/madspace/include/madspace/driver/channel_generator.hpp @@ -49,7 +49,9 @@ class ChannelEventGenerator { bool needs_optimization() const { return (_vegas_optimizer || _discrete_optimizer) && !_status.optimized; } - void set_target_count(double target_count) { _status.count_target = target_count; } + void set_target_count(std::size_t target_count) { + _status.count_target = target_count; + } const std::unordered_set& used_globals() const { return _used_globals; } @@ -61,7 +63,20 @@ class ChannelEventGenerator { void integrate(const GeneratorBatchJob& job); void optimize_vegas(const GeneratorBatchJob& job); double channel_weight_sum(std::size_t event_count); - void start_job(GeneratorBatchJob& job, ResultQueue& result_queue); + void start_job( + GeneratorBatchJob& job, + ResultQueue& result_queue, + std::uint64_t seed, + bool is_survey, + std::size_t survey_pass + ); + // Snapshots max_weight into job.max_weight, at a fixed point independent of + // when submit_unweight_job() actually dispatches it. + void prepare_unweight_job(GeneratorBatchJob& job) const; + // Submits job's unweighting to the thread pool of its own generation context + // (required on GPU: unweighting does a device-to-host copy). + void submit_unweight_job(GeneratorBatchJob& job, ResultQueue& result_queue); + // prepare_unweight_job() + submit_unweight_job() combined. void start_unweight_job(GeneratorBatchJob& job, ResultQueue& result_queue); std::size_t next_vegas_batch_size(); void clear_events(); @@ -127,7 +142,17 @@ class ChannelEventGenerator { std::optional _histogram_function; RunningIntegral _cross_section; double _max_weight = 0.; + // Monotonic per-job counters keying each job's deterministic random stream; + // never reset. Separate for survey/generate so neither depends on the other. + std::size_t _survey_rng_seq = 0; + std::size_t _generate_rng_seq = 0; + // Progress of unweight_file()'s final pass over _weight_file: _unweighted_count + // is the index up to which it's been scanned under the *current* _max_weight, + // and _unweighted_accept_count the number of those accepted. Both reset to 0 + // whenever _max_weight changes (forcing a full rescan), since a changed + // max_weight invalidates every prior accept/reject decision. std::size_t _unweighted_count = 0; + std::size_t _unweighted_accept_count = 0; std::size_t _iters_without_improvement = 0; double _best_rsd = std::numeric_limits::max(); std::vector _large_weights; diff --git a/madspace/include/madspace/driver/context.hpp b/madspace/include/madspace/driver/context.hpp index 13e438084..b17481a1c 100644 --- a/madspace/include/madspace/driver/context.hpp +++ b/madspace/include/madspace/driver/context.hpp @@ -48,35 +48,43 @@ class MatrixElementApi { std::size_t index() const { return _index; } const std::string& file_name() const { return _file_name; } std::vector supported_inputs() const { - bool const* data; int count; + bool const* data; + int count; check_umami_status(_supported_inputs(&data, &count)); std::vector result(UMAMI_INPUT_KEY_COUNT, false); - for (int i = 0; i < count && i < UMAMI_INPUT_KEY_COUNT; ++i) + for (int i = 0; i < count && i < UMAMI_INPUT_KEY_COUNT; ++i) { result[i] = data[i]; + } return result; } std::vector required_inputs() const { - bool const* data; int count; + bool const* data; + int count; check_umami_status(_required_inputs(&data, &count)); std::vector supported = supported_inputs(); for (int i = 0; i < count && i < UMAMI_INPUT_KEY_COUNT; ++i) { if (data[i] && !supported[i]) { - throw_error(std::format( - "input key {} is reported as required but not as supported", i - )); + throw_error( + std::format( + "input key {} is reported as required but not as supported", i + ) + ); } } std::vector result(UMAMI_INPUT_KEY_COUNT, false); - for (int i = 0; i < count && i < UMAMI_INPUT_KEY_COUNT; ++i) + for (int i = 0; i < count && i < UMAMI_INPUT_KEY_COUNT; ++i) { result[i] = data[i]; + } return result; } std::vector supported_outputs() const { - bool const* data; int count; + bool const* data; + int count; check_umami_status(_supported_outputs(&data, &count)); std::vector result(UMAMI_OUTPUT_KEY_COUNT, false); - for (int i = 0; i < count && i < UMAMI_OUTPUT_KEY_COUNT; ++i) + for (int i = 0; i < count && i < UMAMI_OUTPUT_KEY_COUNT; ++i) { result[i] = data[i]; + } return result; } diff --git a/madspace/include/madspace/driver/event_generator.hpp b/madspace/include/madspace/driver/event_generator.hpp index 0e9428dff..0933fbdf1 100644 --- a/madspace/include/madspace/driver/event_generator.hpp +++ b/madspace/include/madspace/driver/event_generator.hpp @@ -3,6 +3,7 @@ #include #include #include +#include #include #include @@ -31,14 +32,18 @@ class EventGenerator { const std::vector& contexts, const std::vector>& channels, const std::string& status_file = "", - const GeneratorConfig& config = default_config + const GeneratorConfig& config = default_config, + std::uint64_t seed = 0 ); - void survey(); + // `survey_pass` salts job seeds so repeated survey() calls on the same + // channel (e.g. re-survey after simplification) don't share a seed stream. + void survey(std::size_t survey_pass = 0); void generate(); void combine_to_compact_npy(const std::string& file_name); void combine_to_lhe_npy(const std::string& file_name, LHECompleter& lhe_completer); void combine_to_lhe( - const std::string& file_name, LHECompleter& lhe_completer, + const std::string& file_name, + LHECompleter& lhe_completer, const LHEMeta& meta = {} ); GeneratorStatus status() const { return _status; } @@ -73,8 +78,42 @@ class EventGenerator { std::vector _channel_optimizing; std::vector _channel_integral_fractions; std::vector _context_job_counts; + // True while a channel has a steady-state batch dispatched but not yet fully + // committed; keeps next_batch_job_count() from double-counting in-flight work. + std::vector _channel_batch_pending; + // generate_deterministic() only: per-channel commit cursor/buffer, analogous + // to _ready_gen/_commit_cursor but ordered per channel instead of globally. + std::vector> _channel_ready_gen; + std::vector _channel_commit_cursor; + std::vector _channel_cursor_set; + // generate_deterministic() only: same as above, for a job's unweight-stage + // completion (tracked separately since it's a distinct completion event). + std::vector> _channel_unweight_ready; + std::vector _channel_unweight_cursor; + // generate_deterministic() only: per-context queue of job ids awaiting + // unweight-stage dispatch, drained with priority by start_jobs(). + std::vector> _context_unweight_queue; ResultQueue _result_queue; + // Base seed for reproducible event generation; 0 means non-deterministic. + std::uint64_t _seed; + + // unweight_all() may run more than once per generate() (a channel's target can + // grow after it looked done, un-finishing it and triggering another round). + // Salted by this counter so repeated calls don't replay the same stream. + std::size_t _unweight_call_index = 0; + + // Scheduling context for the running survey()/generate() call, read by + // start_jobs() to derive job seeds. + bool _survey_job = false; + std::size_t _survey_pass = 0; + + // Deterministic-path state, set from _seed. Generate completions are + // committed in ascending job id, with _commit_cursor as the next id due. + bool _deterministic = false; + std::set _ready_gen; + std::size_t _commit_cursor = 0; + std::chrono::time_point _start_time; std::size_t _start_cpu_microsec; std::chrono::time_point _last_print_time; @@ -84,8 +123,17 @@ class EventGenerator { std::string _status_file; std::unordered_map _timing_data; - bool start_jobs(); + void survey_deterministic(); + void generate_deterministic(); + void commit_generate_job(GeneratorBatchJob& job); + void commit_unweight_job(GeneratorBatchJob& job); + void finish_channel_job(const GeneratorBatchJob& job); + void register_dispatched_ids(std::size_t first_id, std::size_t end_id); + std::size_t next_batch_job_count(std::size_t channel_index) const; + std::size_t start_jobs(); void update_integral(); + void update_integral_status(); + void update_integral_fractions(); void update_counts(); void reset_start_time(); void add_timing_data(const std::string& key); @@ -94,7 +142,8 @@ class EventGenerator { void read_and_combine( std::vector& channel_data, EventBuffer& buffer, - double norm_factor + double norm_factor, + std::mt19937& rand_gen ); void fill_lhe_event( LHECompleter& lhe_completer, diff --git a/madspace/include/madspace/driver/generator_data.hpp b/madspace/include/madspace/driver/generator_data.hpp index bedf6bb2c..d595917dc 100644 --- a/madspace/include/madspace/driver/generator_data.hpp +++ b/madspace/include/madspace/driver/generator_data.hpp @@ -63,6 +63,9 @@ struct GeneratorConfig { int combine_thread_count = -1; double cut_efficiency_threshold = 0.7; std::size_t max_cut_repetitions = 100; + // Cap/floor on a steady-state batch's job count; see next_batch_job_count(). + double generation_batch_fraction = 0.5; + std::size_t min_batch_jobs = 1; }; struct GeneratorStatus { @@ -76,7 +79,7 @@ struct GeneratorStatus { std::size_t count_after_cuts; std::size_t count_after_cuts_opt; double count_unweighted; - double count_target; + std::size_t count_target; std::size_t iterations; bool optimized; bool done; @@ -93,6 +96,7 @@ struct Histogram { struct GeneratorBatchJob { std::size_t channel_index; bool unweight; + // Nonzero: start_jobs() splits this into sub-jobs and dispatches them atomically. std::size_t vegas_batch_size; std::size_t split_job_count; Tensor weights; @@ -104,6 +108,11 @@ struct GeneratorBatchJob { std::size_t context_index; std::size_t job_id; double max_weight; + // Base seed for this job's deterministic random stream; 0 means non-deterministic. + std::uint64_t rng_seed = 0; + // True for a VEGAS-grid-optimization batch, false for a steady-state generation + // batch -- both use vegas_batch_size > 0 for atomic whole-batch dispatch. + bool is_vegas_batch = false; }; void to_json(nlohmann::json& j, const GeneratorStatus& status); diff --git a/madspace/include/madspace/driver/madnis_training.hpp b/madspace/include/madspace/driver/madnis_training.hpp index 932760119..d3a94a8f6 100644 --- a/madspace/include/madspace/driver/madnis_training.hpp +++ b/madspace/include/madspace/driver/madnis_training.hpp @@ -1,6 +1,7 @@ #pragma once #include +#include #include "madspace/compgraphs.hpp" #include "madspace/driver/adam_optimizer.hpp" @@ -41,13 +42,18 @@ class MadnisTraining { double buffer_unweighting_quantile = 0.99; double fixed_cwnet_fraction = 0.33; double softclip_threshold = 0.0; + // Makes generator-job scheduling (CPU only) deterministic given the same + // seed, independent of thread count, at some cost to throughput/memory. + // GPU multi-channel batches are unaffected and stay non-deterministic. + bool reproducible = false; }; MadnisTraining( ContextPtr generator_context, ContextPtr optimizer_context, const Config& config, const std::vector>& integrands, - const std::optional& cwnet + const std::optional& cwnet, + std::uint64_t seed = 0 ); void train_step(std::size_t batch_index); std::vector active_channels() const; @@ -65,6 +71,8 @@ class MadnisTraining { struct SampleJob { SampleBatch samples; SampleBatch unweighted_samples; + // per-channel dispatch sequence, used to commit in order; unused for GPU jobs + std::size_t channel_seq = 0; }; struct ChannelData { std::size_t index; @@ -77,6 +85,13 @@ class MadnisTraining { RuntimePtr generator_runtime = nullptr; RuntimePtr unweighter_runtime = nullptr; SampleBatch buffer; + // commit-ordering state for single-channel generator jobs (see + // process_job_results) + std::size_t next_dispatch_seq = 0; + std::size_t commit_cursor = 0; + std::unordered_map ready_job_ids; + // reproducible mode: samples staged here, flushed into buffer next round + std::vector pending_buffer_samples; }; inline static std::function _abort_check_function = [] {}; @@ -84,6 +99,9 @@ class MadnisTraining { void build_runtimes_and_optimizer(); std::vector compute_channel_sizes(); void start_generator_jobs(const std::vector& channel_fractions); + void maybe_start_generator_jobs( + const std::vector& channel_fractions, bool is_online_attempt + ); TensorVec permute_tensors(const TensorVec& tensors) const; void start_single_job(std::size_t channel_index, std::size_t batch_size); void start_multi_job(const std::vector batch_sizes); @@ -115,6 +133,9 @@ class MadnisTraining { std::vector _arg_permutation; bool _buffer_ready = false; std::vector _active_flavors_count; + std::uint64_t _seed; + // sequence for seeding build_buffered_training_batch's BatchSampler::run() call + std::size_t _buffered_batch_seq = 0; }; class MultiMadnisTraining { @@ -124,7 +145,8 @@ class MultiMadnisTraining { ContextPtr optimizer_context, const MadnisTraining::Config& config, const nested_vector2>& integrands, - const std::vector>& cwnets + const std::vector>& cwnets, + std::uint64_t seed = 0 ); void train(); nested_vector2 active_channels() const; diff --git a/madspace/include/madspace/driver/random.hpp b/madspace/include/madspace/driver/random.hpp new file mode 100644 index 000000000..33a48247f --- /dev/null +++ b/madspace/include/madspace/driver/random.hpp @@ -0,0 +1,77 @@ +#pragma once + +#include +#include +#include +#include + +namespace madspace { + +// Mixes `seed` with `salts` (splitmix64-based) to derive an independent +// sub-seed, e.g. for keeping different random streams derived from the same +// base seed uncorrelated. `seed == 0` is a sentinel for "non-deterministic" +// and passes through unchanged. +std::uint64_t +mix_seed(std::uint64_t seed, std::initializer_list salts = {}); + +// Builds a seeded std::mt19937 from `seed` and `salts`. `seed == 0` seeds +// non-deterministically from std::random_device instead. +std::mt19937 +seeded_rng(std::uint64_t seed, std::initializer_list salts = {}); + +// Deterministic (platform/compiler-independent) 64-bit FNV-1a hash of `s`, +// used to fold a name (e.g. a global's name) into a mix_seed()/seeded_rng() +// salt. Not for anything security-sensitive -- just needs to be stable across +// runs/machines and have a low accidental-collision rate. +std::uint64_t hash_string(std::string_view s); + +// Registry of the top-level salts used across madspace, so independent RNG +// streams derived from the same base seed by different subsystems never +// collide, even if they happen to be given the same base seed. Each is used +// at its documented call site; add new ones here (with a comment pointing at +// the call site and the salt tuple shape it's used in) rather than picking an +// arbitrary local value. +namespace salt { + +// driver/channel_generator.cpp -- ChannelEventGenerator::start_job's per-job +// base seed: mix_seed(seed, {job.channel_index, job_kind_*, ...}). +inline constexpr std::uint64_t job_kind_survey = 1; +inline constexpr std::uint64_t job_kind_generate = 2; + +// driver/channel_generator.cpp -- sub-seed for an individual Runtime::run() +// call within a job: mix_seed(job.rng_seed, {phase_*, call_index}). +inline constexpr std::uint64_t phase_generate = 1; +inline constexpr std::uint64_t phase_unweight = 2; + +// driver/event_generator.cpp -- EventGenerator's post-generation passes: +// seeded_rng(seed, {combine_select}) / seeded_rng(seed, {lhe_complete}) / +// seeded_rng(seed, {lhe_complete, batch_seq}) / +// seeded_rng(seed, {unweight_pass, call_index, channel_index}). +inline constexpr std::uint64_t combine_select = 1; +inline constexpr std::uint64_t lhe_complete = 2; +// Named unweight_pass, not unweight_all, to not collide with the unrelated +// member function EventGenerator::unweight_all(). +inline constexpr std::uint64_t unweight_pass = 3; + +// driver/madnis_training.cpp -- MadnisTraining::start_single_job's per-job +// base seed: mix_seed(seed, {job_kind_madnis_train, channel_index, channel_seq}). +inline constexpr std::uint64_t job_kind_madnis_train = 100; +// build_buffered_training_batch's BatchSampler seed: mix_seed(seed, {..., seq}). +inline constexpr std::uint64_t buffered_batch_sample = 101; +// start_single_job's per-job BufferUnweighter seed: mix_seed(job_seed, {...}). +inline constexpr std::uint64_t madnis_train_unweight = 102; + +// phasespace/*.cpp -- random initialization of a named global (MLP weights +// and biases): see global_init_seed(). +inline constexpr std::uint64_t global_init = 200; + +} // namespace salt + +// Derives the seed used to randomly initialize the global named `name`, +// salted (via salt::kGlobalInit and a hash of `name`) so it collides neither +// between different globals nor with any other subsystem's stream derived +// from the same base `seed`. `seed == 0` passes through unchanged (the usual +// non-deterministic sentinel). +std::uint64_t global_init_seed(std::uint64_t seed, std::string_view name); + +} // namespace madspace diff --git a/madspace/include/madspace/phasespace/channel_weight_network.hpp b/madspace/include/madspace/phasespace/channel_weight_network.hpp index b7e0cbb0e..9c605bf22 100644 --- a/madspace/include/madspace/phasespace/channel_weight_network.hpp +++ b/madspace/include/madspace/phasespace/channel_weight_network.hpp @@ -32,7 +32,8 @@ class ChannelWeightNetwork : public FunctionGenerator { const MLP& mlp() const { return _mlp; } const MomentumPreprocessing& preprocessing() const { return _preprocessing; } - void initialize_globals(ContextPtr context) const; + // seed == 0 (the default) initializes non-deterministically. + void initialize_globals(ContextPtr context, std::uint64_t seed = 0) const; const std::string& mask_name() const { return _mask_name; } private: diff --git a/madspace/include/madspace/phasespace/discrete_flow.hpp b/madspace/include/madspace/phasespace/discrete_flow.hpp index 5dc05a793..de15814dc 100644 --- a/madspace/include/madspace/phasespace/discrete_flow.hpp +++ b/madspace/include/madspace/phasespace/discrete_flow.hpp @@ -18,7 +18,8 @@ class DiscreteFlow : public Mapping { ); const std::vector& option_counts() const { return _option_counts; } std::size_t condition_dim() const { return _condition_dim; } - void initialize_globals(ContextPtr context) const; + // seed == 0 (the default) initializes non-deterministically. + void initialize_globals(ContextPtr context, std::uint64_t seed = 0) const; private: Result build_forward_impl( diff --git a/madspace/include/madspace/phasespace/flow.hpp b/madspace/include/madspace/phasespace/flow.hpp index fcf064084..3e7f772df 100644 --- a/madspace/include/madspace/phasespace/flow.hpp +++ b/madspace/include/madspace/phasespace/flow.hpp @@ -19,8 +19,11 @@ class Flow : public Mapping { ); std::size_t input_dim() const { return _input_dim; } std::size_t condition_dim() const { return _condition_dim; } - void initialize_globals(ContextPtr context) const; - void initialize_from_vegas(ContextPtr context, const std::string& grid_name) const; + // seed == 0 (the default) initializes non-deterministically. + void initialize_globals(ContextPtr context, std::uint64_t seed = 0) const; + void initialize_from_vegas( + ContextPtr context, const std::string& grid_name, std::uint64_t seed = 0 + ) const; private: Result build_forward_impl( diff --git a/madspace/include/madspace/phasespace/mlp.hpp b/madspace/include/madspace/phasespace/mlp.hpp index 45234a8f3..3a9157eba 100644 --- a/madspace/include/madspace/phasespace/mlp.hpp +++ b/madspace/include/madspace/phasespace/mlp.hpp @@ -1,5 +1,7 @@ #pragma once +#include + #include "madspace/driver/context.hpp" #include "madspace/phasespace/base.hpp" @@ -17,7 +19,8 @@ class MLP : public FunctionGenerator { std::size_t input_dim() const { return _input_dim; } std::size_t output_dim() const { return _output_dim; } - void initialize_globals(ContextPtr context) const; + // seed == 0 (the default) initializes non-deterministically. + void initialize_globals(ContextPtr context, std::uint64_t seed = 0) const; std::string last_layer_bias_name() const { return prefixed_name(_prefix, std::format("layer{}.bias", _layers)); } diff --git a/madspace/include/madspace/util.hpp b/madspace/include/madspace/util.hpp index 7b1e43b38..de5d8c9b5 100644 --- a/madspace/include/madspace/util.hpp +++ b/madspace/include/madspace/util.hpp @@ -2,6 +2,7 @@ #include #include +#include #include #include #include diff --git a/madspace/src/cpu/device.hpp b/madspace/src/cpu/device.hpp index a29bb9e5b..26c1a712e 100644 --- a/madspace/src/cpu/device.hpp +++ b/madspace/src/cpu/device.hpp @@ -1,5 +1,9 @@ #pragma once +#include +#include + +#include "madspace/driver/random.hpp" #include "madspace/driver/tensor.hpp" #include "madspace/driver/thread_pool.hpp" #include "simd.hpp" @@ -7,10 +11,17 @@ namespace madspace { namespace cpu { +class CpuRuntime; + class CpuDevice : public Device { public: static constexpr bool is_concurrent = false; + // RNG for op_random/op_random_int/op_unweight. Base implementation falls back + // to the runtime's non-deterministic per-thread stream; SeqCpuDevice overrides + // it with its own seeded generator. + std::mt19937& rand_gen(CpuRuntime& runtime) const; + std::pair allocate(std::size_t size, AllocHint hint) const override { std::pair ret{new std::byte[size], Tensor()}; if (needs_zero_init(hint)) { @@ -67,6 +78,16 @@ class CpuDevice : public Device { CpuDevice() = default; }; +class SeqCpuDevice : public CpuDevice { +public: + explicit SeqCpuDevice(std::uint64_t seed) : _rand_gen(seeded_rng(seed)) {} + + std::mt19937& rand_gen(CpuRuntime&) const { return _rand_gen; } + +private: + mutable std::mt19937 _rand_gen; +}; + class AsyncCpuDevice : public CpuDevice { public: static constexpr bool is_concurrent = true; diff --git a/madspace/src/cpu/runtime.cpp b/madspace/src/cpu/runtime.cpp index 4c7f05474..e83cca19f 100644 --- a/madspace/src/cpu/runtime.cpp +++ b/madspace/src/cpu/runtime.cpp @@ -1,6 +1,7 @@ #include "runtime.hpp" #include +#include #include #include #include @@ -613,10 +614,10 @@ void op_random( auto& runtime = instruction.runtime; device.foreach ( flat_view.shape[0], - [flat_view, &runtime](std::size_t count, std::size_t offset) mutable { + [flat_view, &runtime, &device](std::size_t count, std::size_t offset) mutable { auto output_view = TensorView(flat_view); std::uniform_real_distribution dist; - auto& rand_gen = runtime.rand_gen(); + auto& rand_gen = device.rand_gen(runtime); for (std::size_t i = offset; i < offset + count; ++i) { output_view[i] = dist(rand_gen); } @@ -636,10 +637,12 @@ void op_random_int( auto& runtime = instruction.runtime; device.foreach ( flat_view.shape[0], - [flat_view, max_val, &runtime](std::size_t count, std::size_t offset) mutable { + [flat_view, max_val, &runtime, &device]( + std::size_t count, std::size_t offset + ) mutable { auto output_view = TensorView(flat_view); std::uniform_int_distribution dist(0, max_val - 1); - auto& rand_gen = runtime.rand_gen(); + auto& rand_gen = device.rand_gen(runtime); for (std::size_t i = offset; i < offset + count; ++i) { output_view[i] = dist(rand_gen); } @@ -673,6 +676,7 @@ void op_unweight( indices_view_flat, uw_weights_view_flat, &runtime, + &device, batch_size, &indices, &uw_weights, @@ -683,7 +687,7 @@ void op_unweight( TensorView indices_view(indices_view_flat); TensorView uw_weights_view(uw_weights_view_flat); std::uniform_real_distribution dist; - auto& rand_gen = runtime.rand_gen(); + auto& rand_gen = device.rand_gen(runtime); std::size_t count = 0; for (std::size_t i = 0; i < batch_size; ++i) { double w = weights_view[i], w_max = max_weight_view[i]; @@ -880,6 +884,12 @@ void op_discrete_histogram( } // namespace +std::mt19937& CpuRuntime::rand_gen() { return _rand_gens.get(); } + +std::mt19937& CpuDevice::rand_gen(CpuRuntime& runtime) const { + return runtime.rand_gen(); +} + CpuRuntime::CpuRuntime(const Function& function, ContextPtr context, bool concurrent) : _context(context), _input_count(function.inputs().size()), @@ -904,9 +914,8 @@ CpuRuntime::CpuRuntime(const Function& function, ContextPtr context, bool concur --instr_index; bool is_ready = true; std::size_t dependency_count = 0; - for (auto& out : instr.outputs) { - local_uses_backward.at(out.local_index).push_back(instr_index); - std::size_t source_instr = local_sources_backward.at(out.local_index); + auto check_dependency = [&](std::size_t local_index) { + std::size_t source_instr = local_sources_backward.at(local_index); if (source_instr != -1) { is_ready = false; auto& source_deps = dep_instrs_backward.at(source_instr); @@ -916,15 +925,24 @@ CpuRuntime::CpuRuntime(const Function& function, ContextPtr context, bool concur ++dependency_count; } } + }; + for (auto& out : instr.outputs) { + local_uses_backward.at(out.local_index).push_back(instr_index); + check_dependency(out.local_index); } - dep_counts_backward.at(instr_index) = dependency_count; - if (is_ready) { - _ready_instructions_backward_init.push_back(instr_index); + // also run dependency check for inputs to prevent concurrent writes to + // same gradient + for (auto& in : instr.inputs) { + check_dependency(in.local_index); } - for (auto& in : instr.inputs) { local_sources_backward.at(in.local_index) = instr_index; } + + dep_counts_backward.at(instr_index) = dependency_count; + if (is_ready) { + _ready_instructions_backward_init.push_back(instr_index); + } } _locals_init.resize(function.locals().size()); @@ -1070,22 +1088,36 @@ CpuRuntime::CpuRuntime(const Function& function, ContextPtr context, bool concur } } -TensorVec CpuRuntime::run(const TensorVec& inputs) { +TensorVec CpuRuntime::run(const TensorVec& inputs, std::optional seed) { if (_concurrent && _context->thread_pool().thread_count() > 1) { + if (seed) { + throw std::runtime_error( + "CpuRuntime::run: a seed cannot be used with a concurrently " + "executing runtime" + ); + } auto [outputs, locals, eval_grad] = run_concurrent(inputs, {}, false); return outputs; } else { - return run_single(inputs); + return run_single(inputs, seed); } } std::tuple> CpuRuntime::run_with_grad( - const TensorVec& inputs, const std::vector& input_requires_grad + const TensorVec& inputs, + const std::vector& input_requires_grad, + std::optional seed ) { if (_concurrent && _context->thread_pool().thread_count() > 1) { + if (seed) { + throw std::runtime_error( + "CpuRuntime::run_with_grad: a seed cannot be used with a " + "concurrently executing runtime" + ); + } return run_concurrent(inputs, input_requires_grad, true); } else { - return run_with_grad_single(inputs, input_requires_grad); + return run_with_grad_single(inputs, input_requires_grad, seed); } } @@ -1106,13 +1138,23 @@ std::pair CpuRuntime::run_backward( } } -TensorVec CpuRuntime::run_single(const TensorVec& inputs) const { - auto& device = CpuDevice::instance(); +TensorVec CpuRuntime::run_single( + const TensorVec& inputs, std::optional seed +) const { + if (seed) { + SeqCpuDevice device(*seed); + return run_single_impl(inputs, device); + } + return run_single_impl(inputs, CpuDevice::instance()); +} + +template +TensorVec CpuRuntime::run_single_impl(const TensorVec& inputs, const D& device) const { auto locals = _locals_init; std::copy(inputs.begin(), inputs.end(), locals.begin()); for (auto& instr : _instructions) { - using DeviceType = CpuDevice; + using DeviceType = D; switch (instr.opcode) { case -1: // free memory locals[instr.input_indices[0]].reset(device); @@ -1128,9 +1170,26 @@ TensorVec CpuRuntime::run_single(const TensorVec& inputs) const { } std::tuple> CpuRuntime::run_with_grad_single( - const TensorVec& inputs, const std::vector& input_requires_grad + const TensorVec& inputs, + const std::vector& input_requires_grad, + std::optional seed +) const { + if (seed) { + SeqCpuDevice device(*seed); + return run_with_grad_single_impl(inputs, input_requires_grad, device); + } + return run_with_grad_single_impl( + inputs, input_requires_grad, CpuDevice::instance() + ); +} + +template +std::tuple> +CpuRuntime::run_with_grad_single_impl( + const TensorVec& inputs, + const std::vector& input_requires_grad, + const D& device ) const { - auto& device = CpuDevice::instance(); auto locals = _locals_init; auto requires_grad = _requires_grad_init; std::vector store_local(locals.size()); @@ -1159,7 +1218,7 @@ std::tuple> CpuRuntime::run_with_grad_si } } } - using DeviceType = CpuDevice; + using DeviceType = D; switch (instr.opcode) { case -1: { // free memory auto input_index = instr.input_indices[0]; @@ -1481,3 +1540,16 @@ extern "C" Runtime* build_runtime(const Function& function, ContextPtr context, bool concurrent) { return new CpuRuntime(function, context, concurrent); } + +template TensorVec +CpuRuntime::run_single_impl(const TensorVec&, const CpuDevice&) const; +template TensorVec +CpuRuntime::run_single_impl(const TensorVec&, const SeqCpuDevice&) const; +template std::tuple> +CpuRuntime::run_with_grad_single_impl( + const TensorVec&, const std::vector&, const CpuDevice& +) const; +template std::tuple> +CpuRuntime::run_with_grad_single_impl( + const TensorVec&, const std::vector&, const SeqCpuDevice& +) const; diff --git a/madspace/src/cpu/runtime.hpp b/madspace/src/cpu/runtime.hpp index 9b70fdbfc..a5b283273 100644 --- a/madspace/src/cpu/runtime.hpp +++ b/madspace/src/cpu/runtime.hpp @@ -1,5 +1,6 @@ #pragma once +#include #include #include "madspace/compgraphs/function.hpp" @@ -37,9 +38,13 @@ class CpuRuntime : public Runtime { CpuRuntime(const Function& function, ContextPtr context, bool concurrent); - TensorVec run(const TensorVec& inputs) override; + TensorVec + run(const TensorVec& inputs, + std::optional seed = std::nullopt) override; std::tuple> run_with_grad( - const TensorVec& inputs, const std::vector& input_requires_grad + const TensorVec& inputs, + const std::vector& input_requires_grad, + std::optional seed = std::nullopt ) override; std::pair run_backward( const TensorVec& output_grads, @@ -49,12 +54,25 @@ class CpuRuntime : public Runtime { ) override; Context& context() { return *_context; } - std::mt19937& rand_gen() { return _rand_gens.get(); } + // Fallback non-deterministic per-thread generator, used whenever no per-call + // seed is in play (see cpu::SeqCpuDevice). + std::mt19937& rand_gen(); private: - TensorVec run_single(const TensorVec& inputs) const; + TensorVec + run_single(const TensorVec& inputs, std::optional seed) const; + template + TensorVec run_single_impl(const TensorVec& inputs, const D& device) const; std::tuple> run_with_grad_single( - const TensorVec& inputs, const std::vector& input_requires_grad + const TensorVec& inputs, + const std::vector& input_requires_grad, + std::optional seed + ) const; + template + std::tuple> run_with_grad_single_impl( + const TensorVec& inputs, + const std::vector& input_requires_grad, + const D& device ) const; std::tuple> run_concurrent( const TensorVec& inputs, diff --git a/madspace/src/cpu/tensor.hpp b/madspace/src/cpu/tensor.hpp index d4397ef40..cb0bd7afe 100644 --- a/madspace/src/cpu/tensor.hpp +++ b/madspace/src/cpu/tensor.hpp @@ -305,6 +305,20 @@ inline void tensor_foreach_impl( bool single_job, S... scalar_args ) { + // A broadcast *output* (size(0) != batch_size, e.g. a global's gradient + // accumulator during backward) is shared across the whole batch instead of + // having one slot per batch element. Splitting the batch across worker threads + // would then have multiple threads accumulate into the very same memory + // concurrently with no synchronization, so force a single job in that case + // (the whole batch handled sequentially by one thread). Broadcast *inputs* are + // only ever read, never written, in both forward and backward -- concurrent + // reads of the same memory are safe, so they don't need this. + for (const Tensor* out : outputs) { + if (out->size(0) != batch_size) { + single_job = true; + } + } + // get views to the tensors with the correct types based on the signature of // scalar_func auto flat_views = std::apply( diff --git a/madspace/src/driver/channel_generator.cpp b/madspace/src/driver/channel_generator.cpp index c86da5107..59f182a73 100644 --- a/madspace/src/driver/channel_generator.cpp +++ b/madspace/src/driver/channel_generator.cpp @@ -1,4 +1,6 @@ #include "madspace/driver/channel_generator.hpp" +#include "madspace/driver/random.hpp" +#include "madspace/util.hpp" using namespace madspace; using json = nlohmann::json; @@ -53,7 +55,7 @@ ChannelEventGenerator::ChannelEventGenerator( .count_after_cuts = 0, .count_after_cuts_opt = 0, .count_unweighted = 0., - .count_target = 1., + .count_target = 1, .optimized = false, .done = false }, @@ -177,7 +179,7 @@ ChannelEventGenerator::ChannelEventGenerator( .count_after_cuts = 0, .count_after_cuts_opt = 0, .count_unweighted = 0., - .count_target = 1., + .count_target = 1, .optimized = false, .done = false }, @@ -279,7 +281,7 @@ void ChannelEventGenerator::unweight_file(std::mt19937& rand_gen) { std::size_t buf_size = 1000000; std::uniform_real_distribution rand_dist; EventBuffer buffer(0, 0, weight_file_layout); - std::size_t accept_count = _unweighted_count; + std::size_t accept_count = _unweighted_accept_count; for (std::size_t i = _unweighted_count; i < _weight_file.event_count(); i += buf_size) { _weight_file.seek(i); @@ -296,6 +298,8 @@ void ChannelEventGenerator::unweight_file(std::mt19937& rand_gen) { _weight_file.seek(i); _weight_file.write(buffer); } + _unweighted_count = _weight_file.event_count(); + _unweighted_accept_count = accept_count; _status.count_unweighted = accept_count; } @@ -387,14 +391,41 @@ double ChannelEventGenerator::channel_weight_sum(std::size_t event_count) { return weight_sum; } +// Assigns job's base seed on the scheduling thread, from its logical identity +// only (channel, kind, sequence) -- never from context_index -- so it doesn't +// depend on worker/context assignment. void ChannelEventGenerator::start_job( - GeneratorBatchJob& job, ResultQueue& result_queue + GeneratorBatchJob& job, + ResultQueue& result_queue, + std::uint64_t seed, + bool is_survey, + std::size_t survey_pass ) { + if (seed == 0) { + job.rng_seed = 0; + } else if (is_survey) { + job.rng_seed = mix_seed( + seed, + {job.channel_index, salt::job_kind_survey, survey_pass, _survey_rng_seq++} + ); + } else { + job.rng_seed = mix_seed( + seed, {job.channel_index, salt::job_kind_generate, _generate_rng_seq++} + ); + } _contexts.at(job.context_index) ->thread_pool() .submit([this, &job, &result_queue]() { auto& runtimes = _runtimes.at(job.context_index); auto& context = _contexts.at(job.context_index); + // Each run() call below gets its own sub-seed derived from job's base seed. + std::uint64_t rng_call_index = 0; + auto next_seed = [&]() -> std::optional { + if (job.rng_seed == 0) { + return std::nullopt; + } + return mix_seed(job.rng_seed, {salt::phase_generate, rng_call_index++}); + }; std::size_t max_batch_size = context->device()->device_type() == DeviceType::cpu ? _config.cpu_batch_size @@ -408,8 +439,9 @@ void ChannelEventGenerator::start_job( std::size_t total_count = 0, repetitions = 0; TensorVec all_ps_points; while (true) { - auto ps_points = - runtimes.integrand_channel->run({Tensor({batch_size})}); + auto ps_points = runtimes.integrand_channel->run( + {Tensor({batch_size})}, next_seed() + ); std::size_t acc_count = ps_points .at(_integrand_channel_function.outputs().index_map().at( @@ -444,23 +476,25 @@ void ChannelEventGenerator::start_job( (target_count - total_count) / cut_eff ); } - job.events = runtimes.integrand_common->run(all_ps_points); + job.events = runtimes.integrand_common->run(all_ps_points, next_seed()); job.weights = job.events.at(_field_indices.weight).cpu(); if (runtimes.observable_histograms) { auto hists = runtimes.observable_histograms->run( {job.events.at(_field_indices.weight), - job.events.at(_field_indices.momenta)} + job.events.at(_field_indices.momenta)}, + next_seed() ); for (auto& item : hists) { job.hists.push_back(item.cpu()); } } - if (job.vegas_batch_size != 0) { + if (job.is_vegas_batch) { if (_vegas_optimizer) { auto hist = runtimes.vegas_histogram->run( {job.events.at(_field_indices.random), - job.events.at(_field_indices.weight)} + job.events.at(_field_indices.weight)}, + next_seed() ); for (auto& item : hist) { job.vegas_hist.push_back(item.cpu()); @@ -471,7 +505,7 @@ void ChannelEventGenerator::start_job( job.events.begin() + _field_indices.rest, job.events.end() }; args.push_back(job.events.at(_field_indices.weight)); - auto hist = runtimes.discrete_histogram->run(args); + auto hist = runtimes.discrete_histogram->run(args, next_seed()); for (auto& item : hist) { job.discrete_hist.push_back(item.cpu()); } @@ -482,20 +516,26 @@ void ChannelEventGenerator::start_job( }); } -void ChannelEventGenerator::start_unweight_job( +void ChannelEventGenerator::prepare_unweight_job(GeneratorBatchJob& job) const { + job.max_weight = _max_weight; +} + +void ChannelEventGenerator::submit_unweight_job( GeneratorBatchJob& job, ResultQueue& result_queue ) { - job.max_weight = _max_weight; _contexts.at(job.context_index) ->thread_pool() .submit([this, &job, &result_queue]() { auto& runtimes = _runtimes.at(job.context_index); auto& context = _contexts.at(job.context_index); + std::optional seed = job.rng_seed == 0 + ? std::nullopt + : std::optional(mix_seed(job.rng_seed, {salt::phase_unweight})); std::vector unweighter_args( job.events.begin(), job.events.begin() + _field_indices.random ); unweighter_args.push_back(Tensor(job.max_weight, context->device())); - TensorVec unw_events = runtimes.unweighter->run(unweighter_args); + TensorVec unw_events = runtimes.unweighter->run(unweighter_args, seed); for (auto& item : unw_events) { job.unweighted_events.push_back(item.cpu()); } @@ -504,6 +544,13 @@ void ChannelEventGenerator::start_unweight_job( }); } +void ChannelEventGenerator::start_unweight_job( + GeneratorBatchJob& job, ResultQueue& result_queue +) { + prepare_unweight_job(job); + submit_unweight_job(job, result_queue); +} + std::size_t ChannelEventGenerator::next_vegas_batch_size() { std::size_t batch_size = _batch_size; _batch_size = std::min(_batch_size * 2, _config.max_batch_size); @@ -514,6 +561,7 @@ void ChannelEventGenerator::clear_events() { _status.count_unweighted = 0; _max_weight = 0; _unweighted_count = 0; + _unweighted_accept_count = 0; _status.count_opt = 0; _status.count_after_cuts_opt = 0; _event_file.clear(); @@ -549,8 +597,9 @@ void ChannelEventGenerator::update_max_weight(Tensor weights) { double w_sum = 0; double max_truncation = _config.max_overweight_truncation * - std::min(_status.count_target, - static_cast(_config.freeze_max_weight_after)); + static_cast(std::min( + _status.count_target, _config.freeze_max_weight_after + )); std::size_t count = 0; for (auto w : _large_weights) { if (w < _max_weight) { @@ -563,6 +612,7 @@ void ChannelEventGenerator::update_max_weight(Tensor weights) { _status.count_unweighted *= _max_weight / w; _max_weight = w; _unweighted_count = 0; + _unweighted_accept_count = 0; } break; } diff --git a/madspace/src/driver/event_generator.cpp b/madspace/src/driver/event_generator.cpp index a5e4760a5..52a515444 100644 --- a/madspace/src/driver/event_generator.cpp +++ b/madspace/src/driver/event_generator.cpp @@ -1,11 +1,15 @@ #include "madspace/driver/event_generator.hpp" +#include #include #include #include +#include +#include #include #include "madspace/driver/logger.hpp" +#include "madspace/driver/random.hpp" #include "madspace/util.hpp" using namespace madspace; @@ -16,7 +20,8 @@ EventGenerator::EventGenerator( const std::vector& contexts, const std::vector>& channels, const std::string& status_file, - const GeneratorConfig& config + const GeneratorConfig& config, + std::uint64_t seed ) : _config(config), _status{ @@ -30,7 +35,7 @@ EventGenerator::EventGenerator( .count_after_cuts = 0, .count_after_cuts_opt = 0, .count_unweighted = 0., - .count_target = static_cast(config.target_count), + .count_target = config.target_count, .optimized = false, .done = false }, @@ -41,9 +46,24 @@ EventGenerator::EventGenerator( _channel_optimizing(channels.size()), _channel_integral_fractions(channels.size(), 1.), _context_job_counts(contexts.size()), + _channel_batch_pending(channels.size()), + _channel_ready_gen(channels.size()), + _channel_commit_cursor(channels.size()), + _channel_cursor_set(channels.size()), + _channel_unweight_ready(channels.size()), + _channel_unweight_cursor(channels.size()), + _context_unweight_queue(contexts.size()), + _seed(seed), + _deterministic(seed != 0), _status_file(status_file) {} -void EventGenerator::survey() { +void EventGenerator::survey(std::size_t survey_pass) { + _survey_job = true; + _survey_pass = survey_pass; + if (_deterministic) { + survey_deterministic(); + return; + } reset_start_time(); bool done = false; std::size_t min_iters = _config.survey_min_iters; @@ -80,6 +100,7 @@ void EventGenerator::survey() { .channel_index = i, .unweight = iter >= min_iters - 1, .vegas_batch_size = vegas_batch_size, + .is_vegas_batch = true, }); if (iter >= min_iters) { total_event_count += vegas_batch_size; @@ -136,6 +157,11 @@ void EventGenerator::survey() { } void EventGenerator::generate() { + _survey_job = false; + if (_deterministic) { + generate_deterministic(); + return; + } reset_start_time(); print_gen_init(); @@ -168,6 +194,7 @@ void EventGenerator::generate() { .channel_index = channel_index, .unweight = true, .vegas_batch_size = channel->next_vegas_batch_size(), + .is_vegas_batch = true, }); } } else { @@ -227,14 +254,326 @@ void EventGenerator::generate() { print_gen_update(true); } -bool EventGenerator::start_jobs() { +// Commits a job's generate stage in ascending job id per channel. Doesn't write +// events itself -- if job.unweight, that's deferred to commit_unweight_job(). +void EventGenerator::commit_generate_job(GeneratorBatchJob& job) { + auto& channel = _channels.at(job.channel_index); + auto& channel_job_count = _channel_job_counts.at(job.channel_index); + if (job.is_vegas_batch && channel_job_count == job.split_job_count) { + channel->clear_events(); + } + channel->integrate(job); + update_integral_status(); + channel->update_max_weight(job.weights); + --channel_job_count; + --_context_job_counts.at(job.context_index); + if (job.unweight) { + // Snapshot max_weight here, at this job's fixed commit position. + channel->prepare_unweight_job(job); + _context_unweight_queue.at(job.context_index).push_back(job.job_id); + } + finish_channel_job(job); + print_gen_update(false); +} + +void EventGenerator::commit_unweight_job(GeneratorBatchJob& job) { + auto& channel = _channels.at(job.channel_index); + channel->write_events(job.unweighted_events, job.max_weight); + update_counts(); + --_channel_job_counts.at(job.channel_index); + --_context_job_counts.at(job.context_index); + finish_channel_job(job); + print_gen_update(false); +} + +void EventGenerator::finish_channel_job(const GeneratorBatchJob& job) { + auto& channel_job_count = _channel_job_counts.at(job.channel_index); + if (channel_job_count != 0) { + return; + } + if (job.is_vegas_batch) { + _channels.at(job.channel_index)->optimize_vegas(job); + _channel_optimizing.at(job.channel_index) = false; + } else { + _channel_batch_pending.at(job.channel_index) = false; + } +} + +void EventGenerator::register_dispatched_ids(std::size_t first_id, std::size_t end_id) { + for (std::size_t id = first_id; id < end_id; ++id) { + std::size_t channel_index = _running_jobs.at(id).channel_index; + if (!_channel_cursor_set.at(channel_index)) { + _channel_commit_cursor.at(channel_index) = id; + _channel_unweight_cursor.at(channel_index) = id; + _channel_cursor_set.at(channel_index) = true; + } + } +} + +// Deterministic generate path: jobs still run on the pool, but commit strictly +// in per-channel job-id order, so results are a pure function of the seed. +void EventGenerator::generate_deterministic() { + reset_start_time(); + print_gen_init(); + + while (true) { + _abort_check_function(); + + for (std::size_t channel_index = 0; channel_index < _channels.size(); + ++channel_index) { + auto& channel = _channels.at(channel_index); + double integral_frac = _channel_integral_fractions.at(channel_index); + if (integral_frac > 0 && + channel->status().count_unweighted >= channel->status().count_target) { + continue; + } + if (channel->needs_optimization()) { + if (!_channel_optimizing.at(channel_index)) { + _channel_optimizing.at(channel_index) = true; + _channel_cursor_set.at(channel_index) = false; + _ready_jobs.push_back({ + .channel_index = channel_index, + .unweight = true, + .vegas_batch_size = channel->next_vegas_batch_size(), + .is_vegas_batch = true, + }); + } + } else if (!_channel_batch_pending.at(channel_index)) { + _channel_batch_pending.at(channel_index) = true; + _channel_cursor_set.at(channel_index) = false; + std::size_t batch_jobs = next_batch_job_count(channel_index); + // Single ready_job with vegas_batch_size set so start_jobs() dispatches + // the whole batch atomically, giving the round a contiguous job-id + // range. + _ready_jobs.push_back({ + .channel_index = channel_index, + .unweight = true, + .vegas_batch_size = batch_jobs * _config.cpu_batch_size, + .is_vegas_batch = false, + }); + } + } + + std::size_t job_id_before = _job_id; + std::size_t unweight_dispatched = start_jobs(); + std::size_t round_in_flight = (_job_id - job_id_before) + unweight_dispatched; + register_dispatched_ids(job_id_before, _job_id); + + // Wait for the round to fully commit before resyncing cross-channel fractions. + while (round_in_flight > 0) { + std::size_t job_id = _result_queue.wait(); + --round_in_flight; + auto& job = _running_jobs.at(job_id); + if (job.unweighted_events.size() == 0) { + auto& ready = _channel_ready_gen.at(job.channel_index); + auto& cursor = _channel_commit_cursor.at(job.channel_index); + ready.insert(job_id); + while (ready.erase(cursor) > 0) { + auto& gen_job = _running_jobs.at(cursor); + commit_generate_job(gen_job); + if (!gen_job.unweight) { + _running_jobs.erase(cursor); + } + ++cursor; + } + } else { + auto& ready = _channel_unweight_ready.at(job.channel_index); + auto& cursor = _channel_unweight_cursor.at(job.channel_index); + ready.insert(job_id); + while (ready.erase(cursor) > 0) { + commit_unweight_job(_running_jobs.at(cursor)); + _running_jobs.erase(cursor); + ++cursor; + } + } + std::size_t job_id_refill = _job_id; + std::size_t refill_unweight_dispatched = start_jobs(); + round_in_flight += (_job_id - job_id_refill) + refill_unweight_dispatched; + register_dispatched_ids(job_id_refill, _job_id); + } + + update_integral_fractions(); + + if (_status.done) { + unweight_all(); + } + if (_status.done) { + break; + } + } + print_gen_update(true); +} + +// Deterministic survey path, same commit-ordering approach as generate_deterministic(). +void EventGenerator::survey_deterministic() { + reset_start_time(); + _commit_cursor = _job_id; + _ready_gen.clear(); + bool done = false; + std::size_t min_iters = _config.survey_min_iters; + std::size_t max_iters = std::max(min_iters, _config.survey_max_iters); + double target_precision = _config.survey_target_precision; + + std::size_t total_event_count = 0; + std::size_t done_event_count = 0; + for (auto& channel : _channels) { + std::size_t chan_batch_size = channel->batch_size(); + for (std::size_t iter = channel->status().iterations; iter < min_iters; + ++iter) { + total_event_count += chan_batch_size; + chan_batch_size = std::min(chan_batch_size * 2, _config.max_batch_size); + } + } + print_survey_init(); + + std::size_t iter = 0; + for (; !done && iter < max_iters; ++iter) { + for (std::size_t i = 0; auto channel : _channels) { + if (channel->status().iterations > iter) { + ++i; + continue; + } + if (iter >= min_iters && + channel->cross_section().rel_error() < target_precision) { + ++i; + continue; + } + std::size_t vegas_batch_size = channel->next_vegas_batch_size(); + // Let register_dispatched_ids() re-init the cursor for this new job. + _channel_cursor_set.at(i) = false; + _ready_jobs.push_back({ + .channel_index = i, + .unweight = iter >= min_iters - 1, + .vegas_batch_size = vegas_batch_size, + .is_vegas_batch = true, + }); + if (iter >= min_iters) { + total_event_count += vegas_batch_size; + } + ++i; + } + std::size_t job_id_before = _job_id; + std::size_t unweight_dispatched = start_jobs(); + std::size_t in_flight = (_job_id - job_id_before) + unweight_dispatched; + register_dispatched_ids(job_id_before, _job_id); + done = true; + while (in_flight > 0) { + std::size_t job_id = _result_queue.wait(); + _abort_check_function(); + --in_flight; + auto& job = _running_jobs.at(job_id); + if (job.unweighted_events.size() == 0) { + // Generate-stage completion, globally ordered via + // _ready_gen/_commit_cursor. + _ready_gen.insert(job_id); + while (_ready_gen.erase(_commit_cursor) > 0) { + auto& gen_job = _running_jobs.at(_commit_cursor); + auto& channel = _channels.at(gen_job.channel_index); + auto& channel_job_count = + _channel_job_counts.at(gen_job.channel_index); + if (channel_job_count == gen_job.split_job_count) { + channel->clear_events(); + } + channel->integrate(gen_job); + update_integral(); + channel->update_max_weight(gen_job.weights); + --channel_job_count; + --_context_job_counts.at(gen_job.context_index); + if (gen_job.unweight) { + // Snapshot max_weight at this fixed commit position. + channel->prepare_unweight_job(gen_job); + _context_unweight_queue.at(gen_job.context_index) + .push_back(gen_job.job_id); + } else { + done = false; + if (channel_job_count == 0) { + channel->optimize_vegas(gen_job); + done_event_count += gen_job.vegas_batch_size; + } + _running_jobs.erase(_commit_cursor); + } + ++_commit_cursor; + } + } else { + // Unweight-stage completion, ordered per channel only. + auto& ready = _channel_unweight_ready.at(job.channel_index); + auto& cursor = _channel_unweight_cursor.at(job.channel_index); + ready.insert(job_id); + while (ready.erase(cursor) > 0) { + auto& uw_job = _running_jobs.at(cursor); + auto& channel = _channels.at(uw_job.channel_index); + channel->write_events(uw_job.unweighted_events, uw_job.max_weight); + update_counts(); + auto& channel_job_count = + _channel_job_counts.at(uw_job.channel_index); + --channel_job_count; + --_context_job_counts.at(uw_job.context_index); + if (channel_job_count == 0 && + channel->cross_section().rel_error() < target_precision) { + done = false; + } + if (channel_job_count == 0) { + channel->optimize_vegas(uw_job); + done_event_count += uw_job.vegas_batch_size; + } + _running_jobs.erase(cursor); + ++cursor; + } + } + std::size_t job_id_refill = _job_id; + std::size_t refill_unweight_dispatched = start_jobs(); + in_flight += (_job_id - job_id_refill) + refill_unweight_dispatched; + register_dispatched_ids(job_id_refill, _job_id); + print_survey_update(false, done_event_count, total_event_count, iter); + } + } + print_survey_update(true, done_event_count, total_event_count, iter - 1); +} + +// Sizes the next batch from committed data only, so the estimate never depends +// on jobs still in flight. +std::size_t EventGenerator::next_batch_job_count(std::size_t channel_index) const { + auto& status = _channels.at(channel_index)->status(); + // No data yet: assume a pessimistic efficiency of 1.0. + double efficiency = status.count_opt > 0 + ? status.count_unweighted / static_cast(status.count_opt) + : 1.; + // Floor at 1 event so a near-zero observed efficiency can't blow this up. + double per_job_estimate = + std::max(efficiency * static_cast(_config.cpu_batch_size), 1.); + double remaining = std::max( + static_cast(status.count_target) - status.count_unweighted, 0. + ); + double estimated = _config.generation_batch_fraction * remaining / per_job_estimate; + return std::max( + _config.min_batch_jobs, static_cast(std::ceil(estimated)) + ); +} + +std::size_t EventGenerator::start_jobs() { std::size_t ready_index = 0, context_index = 0; + std::size_t unweight_dispatched = 0; for (auto [context, job_count] : zip(_contexts, _context_job_counts)) { // fill the queue to twice the thread count to keep the worker threads busy std::size_t target_count = 2 * context->thread_pool().thread_count(); std::size_t batch_size = context->device()->device_type() == DeviceType::cpu ? _config.cpu_batch_size : _config.gpu_batch_size; + + // give priority to unweighting jobs to reduce memory usage + auto& unweight_queue = _context_unweight_queue.at(context_index); + std::size_t unweight_index = 0; + for (; job_count < target_count && unweight_index < unweight_queue.size(); + ++unweight_index) { + auto& job = _running_jobs.at(unweight_queue.at(unweight_index)); + _channels.at(job.channel_index)->submit_unweight_job(job, _result_queue); + ++job_count; + ++unweight_dispatched; + } + unweight_queue.erase( + unweight_queue.begin(), unweight_queue.begin() + unweight_index + ); + for (; job_count < target_count && ready_index < _ready_jobs.size(); ++ready_index) { auto ready_job = _ready_jobs.at(ready_index); @@ -247,22 +586,25 @@ bool EventGenerator::start_jobs() { job.split_job_count = split_job_count * (1 + job.unweight); job.job_id = _job_id; job.context_index = context_index; - _channels.at(job.channel_index)->start_job(job, _result_queue); + _channels.at(job.channel_index) + ->start_job(job, _result_queue, _seed, _survey_job, _survey_pass); _channel_job_counts.at(job.channel_index) += 1 + job.unweight; ++_job_id; ++job_count; } } - if (ready_index == _ready_jobs.size()) { - break; - } ++context_index; } _ready_jobs.erase(_ready_jobs.begin(), _ready_jobs.begin() + ready_index); - return ready_index > 0; + return unweight_dispatched; } void EventGenerator::update_integral() { + update_integral_status(); + update_integral_fractions(); +} + +void EventGenerator::update_integral_status() { double total_mean = 0., total_var = 0.; std::size_t total_count = 0, total_count_opt = 0; std::size_t total_count_after_cuts = 0, total_count_after_cuts_opt = 0; @@ -293,19 +635,48 @@ void EventGenerator::update_integral() { _status.count_after_cuts_opt = total_count_after_cuts_opt; _status.iterations = iterations; _status.optimized = optimized; +} + +void EventGenerator::update_integral_fractions() { for (auto [channel, integral_fraction] : zip(_channels, _channel_integral_fractions)) { - integral_fraction = channel->cross_section().mean() / total_mean; - channel->set_target_count(integral_fraction * _config.target_count); + integral_fraction = channel->cross_section().mean() / _status.mean; + } + + // Distribute events between channels, ensure sum is exactly target_count + std::vector counts(_channels.size()); + std::vector remainders(_channels.size()); + std::size_t allocated = 0; + for (auto [count, remainder, fraction] : + zip(counts, remainders, _channel_integral_fractions)) { + double raw = fraction * _config.target_count; + double floored = std::floor(std::isfinite(raw) && raw > 0. ? raw : 0.); + count = static_cast(floored); + remainder = raw - floored; + allocated += count; + } + std::vector order(_channels.size()); + std::iota(order.begin(), order.end(), 0); + std::sort(order.begin(), order.end(), [&](std::size_t a, std::size_t b) { + double rem_a = remainders.at(a), rem_b = remainders.at(b); + return rem_a == rem_b ? a < b : rem_a > rem_b; + }); + std::size_t leftover = + allocated <= _config.target_count ? _config.target_count - allocated : 0; + leftover = std::min(leftover, order.size()); + for (std::size_t index : order | std::views::take(leftover)) { + ++counts.at(index); + } + for (auto [channel, count] : zip(_channels, counts)) { + channel->set_target_count(count); } } void EventGenerator::update_counts() { double total_eff_count = 0.; bool done = true; - for (auto [channel, integral_fraction] : - zip(_channels, _channel_integral_fractions)) { - double chan_target = integral_fraction * _config.target_count; + for (auto& channel : _channels) { + std::size_t chan_target = channel->status().count_target; if (channel->status().count_unweighted < chan_target) { total_eff_count += channel->status().count_unweighted; done = false; @@ -319,6 +690,7 @@ void EventGenerator::update_counts() { void EventGenerator::combine_to_compact_npy(const std::string& file_name) { reset_start_time(); + std::mt19937 select_rng = seeded_rng(_seed, {salt::combine_select}); auto [channel_data, particle_count, norm_factor] = init_combine(); DataLayout layout( EventRecord::layout( @@ -337,7 +709,7 @@ void EventGenerator::combine_to_compact_npy(const std::string& file_name) { print_combine_init(); while (true) { _abort_check_function(); - read_and_combine(channel_data, buffer, norm_factor); + read_and_combine(channel_data, buffer, norm_factor, select_rng); if (buffer.event_count() == 0) { break; } @@ -355,8 +727,9 @@ void EventGenerator::combine_to_lhe_npy( const std::string& file_name, LHECompleter& lhe_completer ) { reset_start_time(); - std::random_device rand_device; - std::mt19937 rand_gen(rand_device()); + std::uint64_t seed = _seed; + std::mt19937 select_rng = seeded_rng(seed, {salt::combine_select}); + std::mt19937 rand_gen = seeded_rng(seed, {salt::lhe_complete}); auto [channel_data, particle_count, norm_factor] = init_combine(); DataLayout in_layout( EventRecord::layout( @@ -388,7 +761,7 @@ void EventGenerator::combine_to_lhe_npy( print_combine_init(); while (true) { _abort_check_function(); - read_and_combine(channel_data, buffer, norm_factor); + read_and_combine(channel_data, buffer, norm_factor, select_rng); if (buffer.event_count() == 0) { break; } @@ -418,11 +791,9 @@ void EventGenerator::combine_to_lhe( const std::string& file_name, LHECompleter& lhe_completer, const LHEMeta& meta ) { reset_start_time(); + std::uint64_t seed = _seed; + std::mt19937 select_rng = seeded_rng(seed, {salt::combine_select}); ThreadPool pool(_config.combine_thread_count); - ThreadResource rand_gens(pool, []() { - std::random_device rand_device; - return std::mt19937(rand_device()); - }); auto [channel_data, particle_count, norm_factor] = init_combine(); std::vector> buffers; std::vector idle_buffers; @@ -444,45 +815,64 @@ void EventGenerator::combine_to_lhe( std::size_t event_count = 0; std::size_t last_update_count = 0; bool done = false; + // Batches are seeded by submission order and written in that same order + // (buffered in ready_slots until their turn), independent of completion timing. + std::size_t next_batch_seq = 0; + std::size_t write_cursor = 0; + std::vector slot_batch_seq(buffers.size()); + std::unordered_map ready_slots; print_combine_init(); while (true) { _abort_check_function(); while (idle_buffers.size() > 0 && !done) { - std::size_t job_id = idle_buffers.back(); - auto& [in_buffer, out_buffer] = buffers.at(job_id); - read_and_combine(channel_data, in_buffer, norm_factor); + std::size_t slot = idle_buffers.back(); + auto& [in_buffer, out_buffer] = buffers.at(slot); + read_and_combine(channel_data, in_buffer, norm_factor, select_rng); if (in_buffer.event_count() == 0) { done = true; break; } idle_buffers.pop_back(); + std::size_t batch_seq = next_batch_seq++; + slot_batch_seq.at(slot) = batch_seq; pool.submit( - [job_id, this, &in_buffer, &out_buffer, &lhe_completer, &rand_gens] { + [slot, batch_seq, seed, this, &in_buffer, &out_buffer, &lhe_completer] { + std::mt19937 rand_gen = seeded_rng( + seed, + {salt::lhe_complete, static_cast(batch_seq)} + ); LHEEvent lhe_event; out_buffer.clear(); for (std::size_t i = 0; i < in_buffer.event_count(); ++i) { fill_lhe_event( - lhe_completer, lhe_event, in_buffer, i, rand_gens.get() + lhe_completer, lhe_event, in_buffer, i, rand_gen ); lhe_event.format_to(out_buffer); } - return job_id; + return slot; } ); } - auto done_jobs = pool.wait_multiple(); - for (std::size_t job_id : done_jobs) { - auto& [in_buffer, out_buffer] = buffers.at(job_id); - idle_buffers.push_back(job_id); + auto done_slots = pool.wait_multiple(); + for (std::size_t slot : done_slots) { + ready_slots.emplace(slot_batch_seq.at(slot), slot); + } + for (auto it = ready_slots.find(write_cursor); it != ready_slots.end(); + it = ready_slots.find(write_cursor)) { + std::size_t slot = it->second; + auto& [in_buffer, out_buffer] = buffers.at(slot); event_file.write_string(out_buffer); event_count += in_buffer.event_count(); + idle_buffers.push_back(slot); + ready_slots.erase(it); + ++write_cursor; if (event_count - last_update_count > 10000) { print_combine_update(event_count); last_update_count = event_count; } } - if (done_jobs.size() == 0 && done) { + if (done_slots.size() == 0 && done) { break; } } @@ -506,15 +896,20 @@ void EventGenerator::add_timing_data(const std::string& key) { } void EventGenerator::unweight_all() { - std::random_device rand_device; - std::mt19937 rand_gen(rand_device()); + std::size_t call_index = _unweight_call_index++; bool done = true; double total_eff_count = 0.; - for (auto [channel, integral_fraction] : - zip(_channels, _channel_integral_fractions)) { + for (std::size_t channel_index = 0; auto& channel : _channels) { + std::mt19937 rand_gen = seeded_rng( + _seed, + {salt::unweight_pass, + static_cast(call_index), + static_cast(channel_index)} + ); channel->unweight_file(rand_gen); + ++channel_index; - double chan_target = integral_fraction * _config.target_count; + std::size_t chan_target = channel->status().count_target; if (channel->status().count_unweighted < chan_target) { total_eff_count += channel->status().count_unweighted; done = false; @@ -572,11 +967,12 @@ EventGenerator::init_combine() { std::size_t count_sum = 0; std::size_t particle_count = 0; double weight_sum = 0.; - for (auto [channel, integral_fraction] : - zip(_channels, _channel_integral_fractions)) { + for (auto& channel : _channels) { particle_count = std::max(particle_count, channel->event_file().particle_count()); - std::size_t count = std::round(integral_fraction * _config.target_count); + // Exact apportioned target (see update_integral_fractions()), so counts + // sum to exactly _config.target_count. + std::size_t count = channel->status().count_target; count_sum += count; channel->event_file().seek(0); weight_sum += channel->channel_weight_sum(count); @@ -596,7 +992,8 @@ EventGenerator::init_combine() { void EventGenerator::read_and_combine( std::vector& channel_data, EventBuffer& buffer, - double norm_factor + double norm_factor, + std::mt19937& rand_gen ) { std::size_t batch_size = 1000; std::size_t event_count = std::min(batch_size, channel_data.back().cum_count); @@ -607,8 +1004,6 @@ void EventGenerator::read_and_combine( bool has_partial = _channels.at(0)->event_layout_extra_flags() & EventRecord::f_partial_weights; - std::random_device rand_device; - std::mt19937 rand_gen(rand_device()); for (std::size_t event_index = 0; event_index < event_count; ++event_index) { std::size_t random_index = std::uniform_int_distribution< std::size_t>(0, channel_data.back().cum_count - 1)(rand_gen); @@ -739,14 +1134,14 @@ void EventGenerator::write_status(const std::string& status, bool force_write) { std::ofstream f(status_tmp_file); nlohmann::json j{ {"status", status}, + {"seed", _seed}, {"process", _status}, {"channels", channel_status()}, {"run_times", _timing_data}, {"histograms", histograms()}, }; f << j.dump(); - // rename atomically deletes the old file and replaces it with the new one - // such that the status file exists at all times + // Atomic rename keeps the status file present at all times. std::filesystem::rename(status_tmp_file, _status_file); } diff --git a/madspace/src/driver/madnis_training.cpp b/madspace/src/driver/madnis_training.cpp index c94d84fb9..2c1b9f2d2 100644 --- a/madspace/src/driver/madnis_training.cpp +++ b/madspace/src/driver/madnis_training.cpp @@ -1,5 +1,8 @@ #include "madspace/driver/madnis_training.hpp" +#include + +#include "madspace/driver/random.hpp" #include "madspace/phasespace/batch_sampler.hpp" using namespace madspace; @@ -9,13 +12,15 @@ MadnisTraining::MadnisTraining( ContextPtr optimizer_context, const Config& config, const std::vector>& integrands, - const std::optional& cwnet + const std::optional& cwnet, + std::uint64_t seed ) : _generator_context(generator_context), _optimizer_context(optimizer_context), _config(config), _channels(integrands.size()), - _cwnet(cwnet) { + _cwnet(cwnet), + _seed(seed) { for (std::size_t index = 0; auto [integrand, channel] : zip(integrands, _channels)) { channel.index = index; @@ -54,7 +59,7 @@ void MadnisTraining::train_step(std::size_t batch_index) { _config.buffer_capacity > 0 && batch_index % (_config.buffered_steps + 1) != 0; TensorVec training_batch; while (true) { - start_generator_jobs(channel_sizes); + maybe_start_generator_jobs(channel_sizes, !try_buffered); if (try_buffered) { if (check_buffered_training_batch(channel_sizes)) { training_batch = build_buffered_training_batch(channel_sizes); @@ -262,9 +267,18 @@ void MadnisTraining::start_generator_jobs(const std::vector& counts if (_running_jobs.size() > 0) { return; } + bool is_gpu = _generator_context->device()->device_type() != DeviceType::cpu; + if (_config.reproducible && !is_gpu) { + // flush buffer samples staged since the last round (see process_job_results) + for (auto& channel : _channels) { + for (auto& pending : channel.pending_buffer_samples) { + buffer_store(channel, pending); + } + channel.pending_buffer_samples.clear(); + } + } _generator_params.copy_from(_optimizer->parameters()); std::size_t chan_count = counts.size(); - bool is_gpu = _generator_context->device()->device_type() != DeviceType::cpu; std::size_t batch_size = is_gpu ? _config.gpu_generator_batch_granularity : _config.cpu_generator_batch_size; @@ -280,7 +294,11 @@ void MadnisTraining::start_generator_jobs(const std::vector& counts ? (target_count - channel.sample_count + batch_size - 1) / batch_size : 0; } - std::size_t available_jobs = _generator_context->thread_pool().thread_count(); + // reproducible mode (CPU only): dispatch the full round in one shot, uncapped + // by thread count, so round contents don't depend on thread count + std::size_t available_jobs = _config.reproducible && !is_gpu + ? std::numeric_limits::max() + : _generator_context->thread_pool().thread_count(); std::vector channel_sizes; std::size_t gpu_subbatches = (_config.gpu_generator_batch_size + _config.gpu_generator_batch_granularity - @@ -349,6 +367,32 @@ void MadnisTraining::start_generator_jobs(const std::vector& counts } } +// Reproducible mode (CPU only) additionally requires an online attempt and the +// online cache to be exhausted or about to be, both deterministic conditions. +void MadnisTraining::maybe_start_generator_jobs( + const std::vector& counts, bool is_online_attempt +) { + if (_running_jobs.size() > 0) { + return; + } + if (_config.reproducible) { + if (!is_online_attempt) { + return; + } + bool depletion_imminent = false; + for (auto [channel, count] : zip(_channels, counts)) { + if (count >= channel.sample_count) { + depletion_imminent = true; + break; + } + } + if (!depletion_imminent) { + return; + } + } + start_generator_jobs(counts); +} + TensorVec MadnisTraining::permute_tensors(const TensorVec& tensors) const { TensorVec ret; ret.reserve(tensors.size()); @@ -358,21 +402,37 @@ TensorVec MadnisTraining::permute_tensors(const TensorVec& tensors) const { return ret; } +// Seed depends only on channel index and channel-local sequence, not on which +// worker thread runs the job. void MadnisTraining::start_single_job( std::size_t channel_index, std::size_t batch_size ) { std::size_t job_id = _job_id; ++_job_id; + auto& channel = _channels.at(channel_index); + std::size_t channel_seq = channel.next_dispatch_seq++; + std::optional seed = _seed == 0 + ? std::nullopt + : std::optional( + mix_seed(_seed, {salt::job_kind_madnis_train, channel_index, channel_seq}) + ); auto& job = std::get<0>(_running_jobs.emplace(job_id, SampleJob{}))->second; + job.channel_seq = channel_seq; _generator_context->thread_pool().submit( - [this, channel_index, batch_size, job_id, &job]() { + [this, channel_index, batch_size, job_id, seed, &job]() { auto& channel = _channels.at(channel_index); - auto samples = channel.generator_runtime->run({Tensor({batch_size})}); + auto samples = channel.generator_runtime->run({Tensor({batch_size})}, seed); job.samples.tensors = permute_tensors(samples); job.samples.size = samples.at(0).size(0); job.samples.channel_index = channel_index; if (channel.unweighter_runtime) { - auto unw_samples = channel.unweighter_runtime->run(samples); + std::optional unweight_seed = seed + ? std::optional( + mix_seed(*seed, {salt::madnis_train_unweight}) + ) + : std::nullopt; + auto unw_samples = + channel.unweighter_runtime->run(samples, unweight_seed); job.unweighted_samples.tensors = permute_tensors(unw_samples); job.unweighted_samples.size = unw_samples.at(0).size(0); job.unweighted_samples.channel_index = channel_index; @@ -382,6 +442,8 @@ void MadnisTraining::start_single_job( ); } +// TODO: GPU multi-channel batches aren't seeded or commit-ordered, unlike +// start_single_job. void MadnisTraining::start_multi_job(const std::vector batch_sizes) { std::size_t job_id = _job_id; ++_job_id; @@ -502,53 +564,77 @@ MadnisTraining::build_buffered_training_batch(const std::vector& counts) } } args.emplace_back(counts); - return _multi_channel_sampler->run(args); + std::optional seed = _seed == 0 + ? std::nullopt + : std::optional( + mix_seed(_seed, {salt::buffered_batch_sample, _buffered_batch_seq++}) + ); + return _multi_channel_sampler->run(args, seed); } void MadnisTraining::process_job_results(const std::vector& job_ids) { for (auto job_id : job_ids) { - auto job = std::move(_running_jobs.extract(job_id).mapped()); + auto& job = _running_jobs.at(job_id); if (job.samples.channel_sizes.size() == 0) { - auto& channel = _channels.at(job.samples.channel_index); - channel.sample_count += job.samples.size; - channel.sample_batches.push_back(std::move(job.samples)); - if (job.unweighted_samples.size > 0) { - buffer_store(channel, job.unweighted_samples); + // single-channel job: mark ready, committed below in dispatch order + _channels.at(job.samples.channel_index) + .ready_job_ids.emplace(job.channel_seq, job_id); + continue; + } + auto multi_job = std::move(_running_jobs.extract(job_id).mapped()); + std::size_t offset = 0, unw_offset = 0, chan_index = 0; + SampleBatch chan_unweighted_samples; + for (auto [channel, chan_size] : + zip(_channels, multi_job.samples.channel_sizes)) { + if (chan_size == 0) { + ++chan_index; + continue; } - } else { - std::size_t offset = 0, unw_offset = 0, chan_index = 0; - SampleBatch chan_unweighted_samples; - for (auto [channel, chan_size] : - zip(_channels, job.samples.channel_sizes)) { - if (chan_size == 0) { - ++chan_index; - continue; - } - channel.sample_count += chan_size; - channel.sample_batches.emplace_back(); - auto& batch = channel.sample_batches.back(); - batch.tensors.reserve(job.samples.tensors.size()); - for (auto& tensor : job.samples.tensors) { - batch.tensors.push_back( - tensor.slice(0, offset, offset + chan_size) + channel.sample_count += chan_size; + channel.sample_batches.emplace_back(); + auto& batch = channel.sample_batches.back(); + batch.tensors.reserve(multi_job.samples.tensors.size()); + for (auto& tensor : multi_job.samples.tensors) { + batch.tensors.push_back(tensor.slice(0, offset, offset + chan_size)); + } + if (multi_job.unweighted_samples.channel_sizes.size() > 0) { + std::size_t unw_chan_size = + multi_job.unweighted_samples.channel_sizes.at(chan_index); + chan_unweighted_samples.tensors.clear(); + chan_unweighted_samples.size = unw_chan_size; + for (auto& tensor : multi_job.unweighted_samples.tensors) { + chan_unweighted_samples.tensors.push_back( + tensor.slice(0, unw_offset, unw_offset + unw_chan_size) ); } - if (job.unweighted_samples.channel_sizes.size() > 0) { - std::size_t unw_chan_size = - job.unweighted_samples.channel_sizes.at(chan_index); - chan_unweighted_samples.tensors.clear(); - chan_unweighted_samples.size = unw_chan_size; - for (auto& tensor : job.unweighted_samples.tensors) { - chan_unweighted_samples.tensors.push_back( - tensor.slice(0, unw_offset, unw_offset + unw_chan_size) - ); - } - buffer_store(channel, chan_unweighted_samples); - unw_offset += unw_chan_size; + buffer_store(channel, chan_unweighted_samples); + unw_offset += unw_chan_size; + } + batch.size = chan_size; + offset += chan_size; + ++chan_index; + } + } + + // commit each channel's ready jobs strictly in dispatch order + for (auto& channel : _channels) { + for (auto it = channel.ready_job_ids.find(channel.commit_cursor); + it != channel.ready_job_ids.end(); + it = channel.ready_job_ids.find(channel.commit_cursor)) { + auto committed_job = std::move(_running_jobs.extract(it->second).mapped()); + channel.ready_job_ids.erase(it); + ++channel.commit_cursor; + channel.sample_count += committed_job.samples.size; + channel.sample_batches.push_back(std::move(committed_job.samples)); + if (committed_job.unweighted_samples.size > 0) { + if (_config.reproducible) { + // flushed into buffer at the start of the next round instead + channel.pending_buffer_samples.push_back( + std::move(committed_job.unweighted_samples) + ); + } else { + buffer_store(channel, committed_job.unweighted_samples); } - batch.size = chan_size; - offset += chan_size; - ++chan_index; } } } @@ -642,7 +728,8 @@ void MadnisTraining::drop_channels() { } std::vector indices(_channels.size()); std::iota(indices.begin(), indices.end(), 0); - std::sort(indices.begin(), indices.end(), [&](auto i, auto j) { + // stable_sort: exact ties (e.g. charge-conjugate channels) must break consistently + std::stable_sort(indices.begin(), indices.end(), [&](auto i, auto j) { return abs_means.at(i) < abs_means.at(j); }); @@ -668,8 +755,7 @@ void MadnisTraining::drop_channels() { return _active_flavors_count.at(flav_index) == 0; } )) { - // cannot drop this channel because one of its flavors is not - // available in any other channel + // a flavor of this channel has no other channel to fall back to continue; } for (std::size_t flav_index : active_flavors) { @@ -719,14 +805,22 @@ MultiMadnisTraining::MultiMadnisTraining( ContextPtr optimizer_context, const MadnisTraining::Config& config, const nested_vector2>& integrands, - const std::vector>& cwnets + const std::vector>& cwnets, + std::uint64_t seed ) : _config(config) { _subprocesses.reserve(integrands.size()); - for (auto [integ, cwnet] : zip(integrands, cwnets)) { + // each subprocess gets its own derived seed so their streams don't collide + for (std::size_t subproc_index = 0; auto [integ, cwnet] : zip(integrands, cwnets)) { _subprocesses.emplace_back( - generator_context, optimizer_context, config, integ, cwnet + generator_context, + optimizer_context, + config, + integ, + cwnet, + seed == 0 ? 0 : mix_seed(seed, {subproc_index}) ); + ++subproc_index; } } diff --git a/madspace/src/driver/random.cpp b/madspace/src/driver/random.cpp new file mode 100644 index 000000000..4ad9b8298 --- /dev/null +++ b/madspace/src/driver/random.cpp @@ -0,0 +1,53 @@ +#include "madspace/driver/random.hpp" + +#include + +namespace madspace { + +std::uint64_t mix_seed(std::uint64_t seed, std::initializer_list salts) { + if (seed == 0) { + return 0; + } + auto splitmix = [](std::uint64_t x) { + x += 0x9E3779B97F4A7C15ull; + x = (x ^ (x >> 30)) * 0xBF58476D1CE4E5B9ull; + x = (x ^ (x >> 27)) * 0x94D049BB133111EBull; + return x ^ (x >> 31); + }; + std::uint64_t h = splitmix(seed); + for (auto s : salts) { + h = splitmix(h ^ splitmix(s)); + } + return h ? h : 1; +} + +std::mt19937 +seeded_rng(std::uint64_t seed, std::initializer_list salts) { + if (seed == 0) { + std::random_device rand_device; + return std::mt19937(rand_device()); + } + std::vector data; + data.reserve(2 + salts.size()); + data.push_back(static_cast(seed & 0xFFFFFFFFu)); + data.push_back(static_cast(seed >> 32)); + data.insert(data.end(), salts.begin(), salts.end()); + std::seed_seq seq(data.begin(), data.end()); + return std::mt19937(seq); +} + +std::uint64_t hash_string(std::string_view s) { + // FNV-1a, 64-bit. + std::uint64_t h = 0xcbf29ce484222325ull; + for (unsigned char c : s) { + h ^= c; + h *= 0x100000001b3ull; + } + return h; +} + +std::uint64_t global_init_seed(std::uint64_t seed, std::string_view name) { + return mix_seed(seed, {salt::global_init, hash_string(name)}); +} + +} // namespace madspace diff --git a/madspace/src/gpu/runtime.cu b/madspace/src/gpu/runtime.cu index 87ff0b0be..0e01bd1e2 100644 --- a/madspace/src/gpu/runtime.cu +++ b/madspace/src/gpu/runtime.cu @@ -1610,11 +1610,14 @@ GpuRuntime::GpuRuntime(const Function& function_arg, ContextPtr context) : ); } -TensorVec GpuRuntime::run(const TensorVec& inputs) { +TensorVec GpuRuntime::run(const TensorVec& inputs, std::optional seed) { auto& gpu_device = *static_cast(_context->device()); auto& streams = _streams.get(); auto& events = _events.get(); gpu_device.activate(); + if (seed) { + check_error(gpurandSetPseudoRandomGeneratorSeed(gpurand_generator(), *seed)); + } auto locals = _locals_init; std::copy(inputs.begin(), inputs.end(), locals.begin()); gpuStream_t main_stream = streams.at(0); @@ -1650,12 +1653,17 @@ TensorVec GpuRuntime::run(const TensorVec& inputs) { } std::tuple> GpuRuntime::run_with_grad( - const TensorVec& inputs, const std::vector& input_requires_grad + const TensorVec& inputs, + const std::vector& input_requires_grad, + std::optional seed ) { auto& gpu_device = *static_cast(_context->device()); auto& streams = _streams.get(); auto& events = _events.get(); gpu_device.activate(); + if (seed) { + check_error(gpurandSetPseudoRandomGeneratorSeed(gpurand_generator(), *seed)); + } auto locals = _locals_init; auto requires_grad = _requires_grad_init; std::vector store_local(locals.size()); diff --git a/madspace/src/gpu/runtime.hpp b/madspace/src/gpu/runtime.hpp index 30930de26..567b28b1b 100644 --- a/madspace/src/gpu/runtime.hpp +++ b/madspace/src/gpu/runtime.hpp @@ -6,6 +6,7 @@ #include "madspace/driver/tensor.hpp" #include +#include namespace madspace { namespace gpu { @@ -31,9 +32,13 @@ class GpuRuntime : public Runtime { }; GpuRuntime(const Function& function, ContextPtr context); - TensorVec run(const TensorVec& inputs) override; + TensorVec + run(const TensorVec& inputs, + std::optional seed = std::nullopt) override; std::tuple> run_with_grad( - const TensorVec& inputs, const std::vector& input_requires_grad + const TensorVec& inputs, + const std::vector& input_requires_grad, + std::optional seed = std::nullopt ) override; std::pair run_backward( const TensorVec& output_grads, diff --git a/madspace/src/phasespace/channel_weight_network.cpp b/madspace/src/phasespace/channel_weight_network.cpp index d84a7a389..972ad7a73 100644 --- a/madspace/src/phasespace/channel_weight_network.cpp +++ b/madspace/src/phasespace/channel_weight_network.cpp @@ -67,8 +67,10 @@ NamedVector ChannelWeightNetwork::build_function_impl( }; } -void ChannelWeightNetwork::initialize_globals(ContextPtr context) const { - _mlp.initialize_globals(context); +void ChannelWeightNetwork::initialize_globals( + ContextPtr context, std::uint64_t seed +) const { + _mlp.initialize_globals(context, seed); context->define_global(_mask_name, DataType::dt_float, {_channel_count}); bool is_cpu = context->device() == cpu_device(); diff --git a/madspace/src/phasespace/discrete_flow.cpp b/madspace/src/phasespace/discrete_flow.cpp index e7119e2e1..e4b9f00a3 100644 --- a/madspace/src/phasespace/discrete_flow.cpp +++ b/madspace/src/phasespace/discrete_flow.cpp @@ -139,13 +139,13 @@ Mapping::Result DiscreteFlow::build_transform( }; } -void DiscreteFlow::initialize_globals(ContextPtr context) const { +void DiscreteFlow::initialize_globals(ContextPtr context, std::uint64_t seed) const { if (_first_prob_name) { initialize_uniform_probs( context, _first_prob_name.value(), _option_counts.at(0) ); } for (auto& subnet : _subnets) { - subnet.initialize_globals(context); + subnet.initialize_globals(context, seed); } } diff --git a/madspace/src/phasespace/flow.cpp b/madspace/src/phasespace/flow.cpp index d3eefd906..682266d0d 100644 --- a/madspace/src/phasespace/flow.cpp +++ b/madspace/src/phasespace/flow.cpp @@ -211,17 +211,17 @@ Flow::Flow( } } -void Flow::initialize_globals(ContextPtr context) const { +void Flow::initialize_globals(ContextPtr context, std::uint64_t seed) const { for (auto& block : _coupling_blocks) { - block.subnet1.initialize_globals(context); - block.subnet2.initialize_globals(context); + block.subnet1.initialize_globals(context, seed); + block.subnet2.initialize_globals(context, seed); } } void Flow::initialize_from_vegas( - ContextPtr context, const std::string& grid_name + ContextPtr context, const std::string& grid_name, std::uint64_t seed ) const { - initialize_globals(context); + initialize_globals(context, seed); auto& last_block = _coupling_blocks.at(_coupling_blocks.size() - 1); vegas_init( context, diff --git a/madspace/src/phasespace/mlp.cpp b/madspace/src/phasespace/mlp.cpp index 920f36594..81542ced7 100644 --- a/madspace/src/phasespace/mlp.cpp +++ b/madspace/src/phasespace/mlp.cpp @@ -3,6 +3,8 @@ #include #include +#include "madspace/driver/random.hpp" + using namespace madspace; namespace { @@ -51,11 +53,9 @@ void initialize_layer( std::size_t output_dim, const std::string& prefix, int layer_index, - std::mt19937& rand_gen, + std::uint64_t seed, bool zeros ) { - double bound = 1 / std::sqrt(input_dim); - std::uniform_real_distribution rand_dist(-bound, bound); auto weight_name = prefixed_name(prefix, std::format("layer{}.weight", layer_index)); auto bias_name = prefixed_name(prefix, std::format("layer{}.bias", layer_index)); @@ -75,14 +75,30 @@ void initialize_layer( } auto weight_view = weight_tensor.view()[0]; - for (std::size_t i = 0; i < output_dim; ++i) { - for (std::size_t j = 0; j < input_dim; ++j) { - weight_view[i][j] = zeros ? 0. : rand_dist(rand_gen); - } - } auto bias_view = bias_tensor.view()[0]; - for (std::size_t i = 0; i < output_dim; ++i) { - bias_view[i] = zeros ? 0. : rand_dist(rand_gen); + if (zeros) { + for (std::size_t i = 0; i < output_dim; ++i) { + for (std::size_t j = 0; j < input_dim; ++j) { + weight_view[i][j] = 0.; + } + bias_view[i] = 0.; + } + } else { + double bound = 1 / std::sqrt(input_dim); + std::uniform_real_distribution rand_dist(-bound, bound); + // Independently seeded per tensor (keyed by its own global name), so + // e.g. reordering layers or adding new ones doesn't shift any other + // tensor's initialized values. + auto weight_rand_gen = seeded_rng(global_init_seed(seed, weight_name)); + for (std::size_t i = 0; i < output_dim; ++i) { + for (std::size_t j = 0; j < input_dim; ++j) { + weight_view[i][j] = rand_dist(weight_rand_gen); + } + } + auto bias_rand_gen = seeded_rng(global_init_seed(seed, bias_name)); + for (std::size_t i = 0; i < output_dim; ++i) { + bias_view[i] = rand_dist(bias_rand_gen); + } } if (!is_cpu) { @@ -133,15 +149,13 @@ MLP::build_function_impl(FunctionBuilder& fb, const NamedVector& args) co }; } -void MLP::initialize_globals(ContextPtr context) const { - std::random_device rand_device; - std::mt19937 rand_gen(rand_device()); +void MLP::initialize_globals(ContextPtr context, std::uint64_t seed) const { std::size_t dim = _input_dim; for (std::size_t i = 1; i < _layers; ++i) { - initialize_layer(context, dim, _hidden_dim, _prefix, i, rand_gen, false); + initialize_layer(context, dim, _hidden_dim, _prefix, i, seed, false); dim = _hidden_dim; } - initialize_layer(context, dim, _output_dim, _prefix, _layers, rand_gen, true); + initialize_layer(context, dim, _output_dim, _prefix, _layers, seed, true); } std::vector MLP::global_names() const { diff --git a/madspace/src/python/madspace.cpp b/madspace/src/python/madspace.cpp index 035b19588..dc57d4285 100644 --- a/madspace/src/python/madspace.cpp +++ b/madspace/src/python/madspace.cpp @@ -842,7 +842,12 @@ PYBIND11_MODULE(_madspace_py, m) { ) .def("input_dim", &MLP::input_dim) .def("output_dim", &MLP::output_dim) - .def("initialize_globals", &MLP::initialize_globals, py::arg("context")); + .def( + "initialize_globals", + &MLP::initialize_globals, + py::arg("context"), + py::arg("seed") = 0 + ); py::classh(m, "Flow") .def( @@ -866,12 +871,18 @@ PYBIND11_MODULE(_madspace_py, m) { ) .def("input_dim", &Flow::input_dim) .def("condition_dim", &Flow::condition_dim) - .def("initialize_globals", &Flow::initialize_globals, py::arg("context")) + .def( + "initialize_globals", + &Flow::initialize_globals, + py::arg("context"), + py::arg("seed") = 0 + ) .def( "initialize_from_vegas", &Flow::initialize_from_vegas, py::arg("context"), - py::arg("grid_name") + py::arg("grid_name"), + py::arg("seed") = 0 ); py::classh( @@ -927,7 +938,8 @@ PYBIND11_MODULE(_madspace_py, m) { .def( "initialize_globals", &ChannelWeightNetwork::initialize_globals, - py::arg("context") + py::arg("context"), + py::arg("seed") = 0 ); py::classh(m, "DiscreteHistogram") @@ -972,7 +984,10 @@ PYBIND11_MODULE(_madspace_py, m) { .def("option_counts", &DiscreteFlow::option_counts) .def("condition_dim", &DiscreteFlow::condition_dim) .def( - "initialize_globals", &DiscreteFlow::initialize_globals, py::arg("context") + "initialize_globals", + &DiscreteFlow::initialize_globals, + py::arg("context"), + py::arg("seed") = 0 ); py::classh(m, "VegasGridOptimizer") @@ -1363,7 +1378,8 @@ PYBIND11_MODULE(_madspace_py, m) { ) .def_readwrite( "softclip_threshold", &MadnisTraining::Config::softclip_threshold - ); + ) + .def_readwrite("reproducible", &MadnisTraining::Config::reproducible); py::classh(m, "MadnisTraining") .def( @@ -1372,12 +1388,14 @@ PYBIND11_MODULE(_madspace_py, m) { ContextPtr, const MadnisTraining::Config&, const std::vector>&, - const std::optional&>(), + const std::optional&, + std::uint64_t>(), py::arg("generator_context"), py::arg("optimizer_context"), py::arg("config"), py::arg("integrands"), - py::arg("cwnet") + py::arg("cwnet"), + py::arg("seed") = 0 ) .def("train_step", &MadnisTraining::train_step, py::arg("batch_index")) .def("active_channels", &MadnisTraining::active_channels) @@ -1390,12 +1408,14 @@ PYBIND11_MODULE(_madspace_py, m) { ContextPtr, const MadnisTraining::Config&, const nested_vector2>&, - const std::vector>&>(), + const std::vector>&, + std::uint64_t>(), py::arg("generator_context"), py::arg("optimizer_context"), py::arg("config"), py::arg("integrands"), - py::arg("cwnets") + py::arg("cwnets"), + py::arg("seed") = 0 ) .def("train", &MultiMadnisTraining::train) .def("active_channels", &MultiMadnisTraining::active_channels); @@ -1429,7 +1449,11 @@ PYBIND11_MODULE(_madspace_py, m) { .def_readwrite( "cut_efficiency_threshold", &GeneratorConfig::cut_efficiency_threshold ) - .def_readwrite("max_cut_repetitions", &GeneratorConfig::max_cut_repetitions); + .def_readwrite("max_cut_repetitions", &GeneratorConfig::max_cut_repetitions) + .def_readwrite( + "generation_batch_fraction", &GeneratorConfig::generation_batch_fraction + ) + .def_readwrite("min_batch_jobs", &GeneratorConfig::min_batch_jobs); py::classh(m, "GeneratorStatus") .def(py::init<>()) @@ -1712,7 +1736,8 @@ PYBIND11_MODULE(_madspace_py, m) { const std::vector&, const std::vector>&, const std::string&, - const GeneratorConfig&>(), + const GeneratorConfig&, + std::uint64_t>(), py::arg("contexts"), py::arg("channels"), py::arg("status_file") = "", @@ -1720,9 +1745,10 @@ PYBIND11_MODULE(_madspace_py, m) { "config", EventGenerator::default_config, "EventGenerator.default_config" - ) + ), + py::arg("seed") = 0 ) - .def("survey", &EventGenerator::survey) + .def("survey", &EventGenerator::survey, py::arg("survey_pass") = 0) .def("generate", &EventGenerator::generate) .def( "combine_to_compact_npy", diff --git a/tests/acceptance_tests/test_mg7_reproducibility.py b/tests/acceptance_tests/test_mg7_reproducibility.py new file mode 100644 index 000000000..ea91a3c52 --- /dev/null +++ b/tests/acceptance_tests/test_mg7_reproducibility.py @@ -0,0 +1,425 @@ +################################################################################ +# +# Copyright (c) 2009 The MadGraph5_aMC@NLO Development team and Contributors +# +# This file is a part of the MadGraph5_aMC@NLO project, an application which +# automatically generates Feynman diagrams and matrix elements for arbitrary +# high-energy processes in the Standard Model and beyond. +# +# It is subject to the MadGraph5_aMC@NLO license which should accompany this +# distribution. +# +# For more information, visit madgraph.phys.ucl.ac.be and amcatnlo.web.cern.ch +# +################################################################################ +"""MG7 (madspace) reproducibility tests. + +Guards the seeded event-generation path: given the same run_card ``seed``, +independent mg7 runs must produce byte-identical LHE output (hashed with +sha256); a different seed must change the result. Covers: + + * ``test_vegas_reproducibility_mg7`` -- a plain survey (VEGAS-optimized) + + generate run, through ``bin/generate_events``. + * ``test_madnis_reproducibility_mg7`` -- a full survey + madnis training + + generate run, through ``bin/generate_events`` (no gridpack export), with + ``madnis.reproducible`` left at its default (off). Unlike the gridpack + test below, this exercises the training phase's own seeding + (single-channel CPU sample generation only); training's own round + scheduling is not required to be deterministic in this mode, so this test + may occasionally be flaky (see test_madnis_reproducible_mode_mg7). + * ``test_madnis_reproducible_mode_mg7`` -- same, but with + ``madnis.reproducible = True``: training is required to be + thread-count-independent and byte-identical for the same seed, both with + online-only training and with buffered (off-policy replay) training + enabled. GPU multi-channel batches are still not covered by this mode. + * ``test_gridpack_reproducibility_mg7`` -- a gridpack trained with madnis, + then standalone event generation from that gridpack (its own + ``bin/generate_events --seed``), which is the normal way a gridpack is + used and does not re-run survey/training. + * ``test_gridpack_reproducibility_vegas_mg7`` -- same as above, but for a + plain VEGAS-optimized gridpack (no madnis training). + +Process: ``p p > t t~`` (hadronic, needs the NNPDF23_lo_as_0130_qed PDF grid; +self-skips if the mg7 runtime stack / PDF is unavailable, same as +test_check_xsec_processes_mg7.py). + +Run locally with e.g.:: + + ./tests/test_manager.py test_.*_reproducibility_mg7 -pA -t0 -l INFO + +One knob is read from the environment so the CI can dial it without touching +the code: + + * ``MG7_REPRO_EVENTS`` -- events per run (default 10000). +""" + +from __future__ import absolute_import +from __future__ import division + +import glob +import hashlib +import os +import shutil +import subprocess +import sys +import tempfile +import unittest + +import madgraph.interface.master_interface as MGCmd +from madgraph.various.banner import RunCardMG7 + +pjoin = os.path.join + +_PROCESS = 'p p > t t~' +_EVENTS = int(os.environ.get('MG7_REPRO_EVENTS', 10000)) +_SEED_A = 424242 +_SEED_B = 909090 + + +def _mg7_datadir_or_skip(test): + """Return an LHAPDF data dir that contains the NNPDF23_lo_as_0130_qed set, + or ``skipTest`` (on *test*) when the mg7 runtime stack (madspace + LHAPDF + + the run_card.toml default PDF) is unavailable.""" + try: + import madspace + has_mg7 = hasattr(madspace, 'ChannelEventGenerator') + except ImportError: + has_mg7 = False + if not has_mg7: + test.skipTest('mg7 runtime stack (madspace) unavailable') + + candidates = [] + if os.environ.get('LHAPDF_DATA_PATH'): + candidates.extend(os.environ['LHAPDF_DATA_PATH'].split(os.pathsep)) + try: + out = subprocess.check_output(['lhapdf-config', '--datadir'], + stderr=subprocess.DEVNULL).decode().strip() + if out: + candidates.append(out) + except Exception: + pass + for d in candidates: + if d and os.path.isdir(d) and glob.glob(pjoin(d, 'NNPDF23_lo_as_0130_qed*')): + return d + test.skipTest('NNPDF23_lo_as_0130_qed LHAPDF data not found ' + '(set $LHAPDF_DATA_PATH)') + + +def _lhe_hash(path): + """sha256 of an LHE file's physics content: everything from ```` + onward, i.e. excluding the ``
`` block. The header embeds the + run_card.toml/param_card/proc_card verbatim, so it legitimately differs + between runs that vary settings unrelated to the physics content (e.g. + the seed itself, or the cpu thread pool size) -- only the generated + events (and the cross-section info, itself a deterministic + function of the seed) should drive this hash.""" + with open(path, 'rb') as f: + content = f.read() + marker = b'
\n' + index = content.find(marker) + if index == -1: + raise AssertionError('no found in %s' % path) + return hashlib.sha256(content[index + len(marker):]).hexdigest() + + +def _find_lhe(run_path): + """Locate the LHE file produced under Events//, or fail.""" + matches = sorted(glob.glob(pjoin(run_path, 'Events', '*', 'events.lhe'))) + if not matches: + raise AssertionError('no events.lhe produced under %s' % run_path) + return matches[-1] + + +def _run(cmd, cwd, log_path, env, what): + """Run *cmd*, raising with the log tail on a non-zero exit.""" + with open(log_path, 'w') as logfh: + ret = subprocess.call(cmd, cwd=cwd, env=env, stdout=logfh, + stderr=subprocess.STDOUT) + if ret != 0: + with open(log_path) as f: + tail = ''.join(f.readlines()[-60:]) + raise AssertionError('%s failed (exit %d, see %s)\n\n%s' + % (what, ret, log_path, tail)) + + +class MG7ReproducibilityTest(unittest.TestCase): + """LHE-hash reproducibility of mg7 (madspace) event generation for a + fixed run_card seed, and non-reproducibility across distinct seeds.""" + + def setUp(self): + self.path = tempfile.mkdtemp(prefix='mg7_repro_') + + def tearDown(self): + shutil.rmtree(self.path, ignore_errors=True) + + def _output_process(self, run_name): + run_dir = pjoin(self.path, run_name) + mg = MGCmd.MasterCmd() + mg.no_notification() + mg.exec_cmd('set automatic_html_opening False --no_save') + mg.exec_cmd('generate %s' % _PROCESS) + mg.exec_cmd('output mg7 %s' % run_dir) + return run_dir + + def _set_run_card(self, toml_path, **settings): + """settings keys are 'section.key' RunCardMG7 names.""" + rc = RunCardMG7(toml_path) + for key, value in settings.items(): + rc.set(key, value, user=True) + rc.write(toml_path) + + def _generate_and_hash(self, run_dir, datadir, seed, tag, **run_card_settings): + """Set run_card seed (plus any extra 'section.key' settings), run + bin/generate_events -f from scratch, and return the sha256 of the + resulting LHE file.""" + shutil.rmtree(pjoin(run_dir, 'Events'), ignore_errors=True) + toml = pjoin(run_dir, 'Cards', 'run_card.toml') + self._set_run_card(toml, **{'run.seed': seed}, **run_card_settings) + env = dict(os.environ) + env['LHAPDF_DATA_PATH'] = datadir + _run([sys.executable, pjoin(run_dir, 'bin', 'generate_events'), '-f'], + run_dir, pjoin(run_dir, 'gen_%s.log' % tag), env, + 'mg7 generate_events') + return _lhe_hash(_find_lhe(run_dir)) + + def _generate_and_hash_gridpack(self, gridpack_dir, datadir, seed, tag): + """Run the gridpack's own bin/generate_events --seed from scratch, and + return the sha256 of the resulting LHE file.""" + shutil.rmtree(pjoin(gridpack_dir, 'Events'), ignore_errors=True) + env = dict(os.environ) + env['LHAPDF_DATA_PATH'] = datadir + _run([sys.executable, pjoin(gridpack_dir, 'bin', 'generate_events'), + '--seed', str(seed), '--events', str(_EVENTS), + '--output_format', 'lhe'], + gridpack_dir, pjoin(gridpack_dir, 'gen_%s.log' % tag), env, + 'gridpack generate_events') + return _lhe_hash(_find_lhe(gridpack_dir)) + + def test_vegas_reproducibility_mg7(self): + """Two VEGAS-optimized runs with the same seed produce byte-identical + LHE files, including when the cpu thread pool is shrunk to 1; a + different seed changes the result.""" + datadir = _mg7_datadir_or_skip(self) + run_dir = self._output_process('vegas') + toml = pjoin(run_dir, 'Cards', 'run_card.toml') + self._set_run_card( + toml, + **{ + 'generation.events': _EVENTS, + # Keep the test scoped to madspace's own seeding: LHE-level + # post-processing (systematics) is orthogonal and would only + # add runtime here. + 'postprocessing.systematics': False, + } + ) + + # Every call sets 'run.cpu_thread_pool_size' explicitly so it never + # leaks from a previous call via the run_card left on disk. + hash_a1 = self._generate_and_hash( + run_dir, datadir, _SEED_A, 'a1', **{'run.cpu_thread_pool_size': -1}) + hash_a2 = self._generate_and_hash( + run_dir, datadir, _SEED_A, 'a2', **{'run.cpu_thread_pool_size': -1}) + self.assertEqual(hash_a1, hash_a2, + 'same seed produced different LHE content') + + hash_a_single = self._generate_and_hash( + run_dir, datadir, _SEED_A, 'a_single', + **{'run.cpu_thread_pool_size': 1}) + self.assertEqual(hash_a1, hash_a_single, + 'same seed produced different LHE content with a ' + 'single-threaded cpu thread pool') + + hash_b = self._generate_and_hash( + run_dir, datadir, _SEED_B, 'b', **{'run.cpu_thread_pool_size': -1}) + self.assertNotEqual(hash_a1, hash_b, + 'different seeds produced identical LHE content') + + def test_madnis_reproducibility_mg7(self): + """Two full survey + madnis training + generate runs (no gridpack + export) with the same seed produce byte-identical LHE files; a + different seed changes the result. Unlike + test_gridpack_reproducibility_mg7, training is re-run from scratch + each time here (not frozen into a gridpack first), so this is the + test that actually exercises MultiMadnisTraining's own seeding + (single-channel CPU sample generation, see salt::job_kind_madnis_train + in random.hpp). madnis.reproducible is left at its default (off) here, + so -- per the known limitation documented at + MadnisTraining::start_generator_jobs -- this is not expected to be + thread-count independent and may occasionally be flaky; see + test_madnis_reproducible_mode_mg7 for the mode that guarantees it.""" + datadir = _mg7_datadir_or_skip(self) + run_dir = self._output_process('madnis') + toml = pjoin(run_dir, 'Cards', 'run_card.toml') + self._set_run_card( + toml, + **{ + 'generation.events': _EVENTS, + 'postprocessing.systematics': False, + 'madnis.enable': True, + 'madnis.train_batches': 100, + } + ) + + # Every call sets 'run.cpu_thread_pool_size' explicitly so it never + # leaks from a previous call via the run_card left on disk. + hash_a1 = self._generate_and_hash( + run_dir, datadir, _SEED_A, 'a1', **{'run.cpu_thread_pool_size': -1}) + hash_a2 = self._generate_and_hash( + run_dir, datadir, _SEED_A, 'a2', **{'run.cpu_thread_pool_size': -1}) + self.assertEqual(hash_a1, hash_a2, + 'same seed produced different LHE content') + + hash_b = self._generate_and_hash( + run_dir, datadir, _SEED_B, 'b', **{'run.cpu_thread_pool_size': -1}) + self.assertNotEqual(hash_a1, hash_b, + 'different seeds produced identical LHE content') + + def test_madnis_reproducible_mode_mg7(self): + """With madnis.reproducible = True (CPU codepath), madnis training + itself -- not just event generation -- is byte-identical for the same + seed independent of the cpu thread pool size, both with online-only + training and with buffered (off-policy replay) training enabled. + Covers the scheduling fix in MadnisTraining::maybe_start_generator_jobs + / start_generator_jobs (deterministic round dispatch + deferred buffer + flush) and the seeding of BufferUnweighter/BatchSampler (see + salt::madnis_train_unweight / salt::buffered_batch_sample in + random.hpp) -- both required for this test to pass reliably.""" + datadir = _mg7_datadir_or_skip(self) + run_dir = self._output_process('madnis_reproducible') + toml = pjoin(run_dir, 'Cards', 'run_card.toml') + base_settings = { + 'generation.events': _EVENTS, + 'postprocessing.systematics': False, + 'madnis.enable': True, + 'madnis.train_batches': 60, + 'madnis.reproducible': True, + } + + # Online-only training (buffering disabled). + self._set_run_card(toml, **base_settings, **{'madnis.buffer_capacity': 0}) + hash_online_multi = self._generate_and_hash( + run_dir, datadir, _SEED_A, 'online_multi', + **{'run.cpu_thread_pool_size': -1}) + hash_online_single = self._generate_and_hash( + run_dir, datadir, _SEED_A, 'online_single', + **{'run.cpu_thread_pool_size': 1}) + self.assertEqual( + hash_online_multi, hash_online_single, + 'madnis.reproducible training (buffering disabled) produced ' + 'different LHE content with a single-threaded cpu thread pool') + + # Buffered (off-policy replay) training enabled. + self._set_run_card( + toml, + **base_settings, + **{ + 'madnis.buffer_capacity': 3000, + 'madnis.buffered_steps': 3, + 'madnis.minimum_buffer_size': 500, + } + ) + hash_buffered_multi = self._generate_and_hash( + run_dir, datadir, _SEED_A, 'buffered_multi', + **{'run.cpu_thread_pool_size': -1}) + hash_buffered_single = self._generate_and_hash( + run_dir, datadir, _SEED_A, 'buffered_single', + **{'run.cpu_thread_pool_size': 1}) + self.assertEqual( + hash_buffered_multi, hash_buffered_single, + 'madnis.reproducible training (buffering enabled) produced ' + 'different LHE content with a single-threaded cpu thread pool') + + hash_buffered_b = self._generate_and_hash( + run_dir, datadir, _SEED_B, 'buffered_b', + **{'run.cpu_thread_pool_size': -1}) + self.assertNotEqual( + hash_buffered_multi, hash_buffered_b, + 'different seeds produced identical LHE content') + + def test_gridpack_reproducibility_mg7(self): + """A gridpack trained with madnis: two standalone event-generation + runs from that gridpack with the same seed produce byte-identical LHE + files; a different seed changes the result. Training itself is not + required to be deterministic -- only the trained gridpack's event + generation is under test here.""" + datadir = _mg7_datadir_or_skip(self) + run_dir = self._output_process('gridpack') + toml = pjoin(run_dir, 'Cards', 'run_card.toml') + self._set_run_card( + toml, + **{ + 'generation.events': _EVENTS, + 'postprocessing.systematics': False, + 'gridpack.save_gridpack': True, + 'madnis.enable': True, + 'madnis.train_batches': 100, + } + ) + + env = dict(os.environ) + env['LHAPDF_DATA_PATH'] = datadir + _run([sys.executable, pjoin(run_dir, 'bin', 'generate_events'), '-f'], + run_dir, pjoin(run_dir, 'train.log'), env, + 'mg7 madnis training run') + + gridpack_dir = pjoin(run_dir, 'Events', 'run_01', 'gridpack') + self.assertTrue(os.path.isdir(gridpack_dir), + 'gridpack was not produced under %s' % run_dir) + + hash_a1 = self._generate_and_hash_gridpack( + gridpack_dir, datadir, _SEED_A, 'a1') + hash_a2 = self._generate_and_hash_gridpack( + gridpack_dir, datadir, _SEED_A, 'a2') + self.assertEqual(hash_a1, hash_a2, + 'same seed produced different LHE content') + + hash_b = self._generate_and_hash_gridpack( + gridpack_dir, datadir, _SEED_B, 'b') + self.assertNotEqual(hash_a1, hash_b, + 'different seeds produced identical LHE content') + + def test_gridpack_reproducibility_vegas_mg7(self): + """A plain VEGAS-optimized gridpack (no madnis training): two + standalone event-generation runs from that gridpack with the same + seed produce byte-identical LHE files; a different seed changes the + result. Companion to test_gridpack_reproducibility_mg7, which covers + the madnis-trained case; the survey/optimization itself is not + required to be deterministic here either -- only the gridpack's own + event generation is under test.""" + datadir = _mg7_datadir_or_skip(self) + run_dir = self._output_process('gridpack_vegas') + toml = pjoin(run_dir, 'Cards', 'run_card.toml') + self._set_run_card( + toml, + **{ + 'generation.events': _EVENTS, + 'postprocessing.systematics': False, + 'gridpack.save_gridpack': True, + } + ) + + env = dict(os.environ) + env['LHAPDF_DATA_PATH'] = datadir + _run([sys.executable, pjoin(run_dir, 'bin', 'generate_events'), '-f'], + run_dir, pjoin(run_dir, 'survey.log'), env, + 'mg7 vegas survey run') + + gridpack_dir = pjoin(run_dir, 'Events', 'run_01', 'gridpack') + self.assertTrue(os.path.isdir(gridpack_dir), + 'gridpack was not produced under %s' % run_dir) + + hash_a1 = self._generate_and_hash_gridpack( + gridpack_dir, datadir, _SEED_A, 'a1') + hash_a2 = self._generate_and_hash_gridpack( + gridpack_dir, datadir, _SEED_A, 'a2') + self.assertEqual(hash_a1, hash_a2, + 'same seed produced different LHE content') + + hash_b = self._generate_and_hash_gridpack( + gridpack_dir, datadir, _SEED_B, 'b') + self.assertNotEqual(hash_a1, hash_b, + 'different seeds produced identical LHE content') + + +if __name__ == '__main__': + unittest.main()