Skip to content

updates with GB10 tests #2

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
sdvillal merged 5 commits into from
Feb 12, 2026
Merged

updates with GB10 tests #2

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
232200b
updates with GB10 tests
jandom Feb 10, 2026
9687822
cleanup
jandom Feb 10, 2026
fa1508d
harmonize
jandom Feb 10, 2026
ecc84cd
linting data_module.py
jandom Feb 10, 2026
534b8af
speculative changes
jandom Feb 10, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
24 changes: 11 additions & 13 deletions openfold3/core/data/framework/data_module.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -162,8 +162,7 @@ class DataModuleConfig(BaseModel):
# Custom worker init function with manual data seed
# (top level to enable pickling working regardless of forking strategy and platform)
def _worker_init_function_with_data_seed(
self,
worker_id: int, rank: int | None = None
self, worker_id: int, rank: int | None = None
) -> None:
"""Modified default Lightning worker_init_fn with manual data seed.

Expand All @@ -182,19 +181,13 @@ def _worker_init_function_with_data_seed(
process_seed = self.data_seed
# back out the base seed so we can use all the bits
base_seed = process_seed - worker_id
seed_sequence = _generate_seed_sequence(
base_seed, worker_id, global_rank, count=4
)
seed_sequence = _generate_seed_sequence(base_seed, worker_id, global_rank, count=4)
torch.manual_seed(seed_sequence[0]) # torch takes a 64-bit seed
random.seed(
(seed_sequence[1] << 32) | seed_sequence[2]
) # combine two 64-bit seeds
random.seed((seed_sequence[1] << 32) | seed_sequence[2]) # combine two 64-bit seeds
if _NUMPY_AVAILABLE:
import numpy as np

np.random.seed(
seed_sequence[3] & 0xFFFFFFFF
) # numpy takes 32-bit seed only
np.random.seed(seed_sequence[3] & 0xFFFFFFFF) # numpy takes 32-bit seed only


class DataModule(pl.LightningDataModule):
Expand Down Expand Up @@ -226,7 +219,10 @@ def _initialize_next_dataset_indices(self):

def setup(self, stage=None):
from functools import partial
self.worker_init_function_with_data_seed = partial(_worker_init_function_with_data_seed, self)

self.worker_init_function_with_data_seed = partial(
_worker_init_function_with_data_seed, self
)
self.generator = torch.Generator(device="cpu").manual_seed(self.data_seed)

self.datasets_by_mode = {k: [] for k in DatasetMode}
Expand Down Expand Up @@ -447,7 +443,9 @@ def generate_dataloader(self, mode: DatasetMode):
generator=self.generator,
worker_init_fn=self.worker_init_function_with_data_seed,
# https://github.com/pytorch/pytorch/issues/87688
multiprocessing_context = "fork" if torch.backends.mps.is_available() and num_workers else None
multiprocessing_context="fork"
if torch.backends.mps.is_available() and num_workers
else None,
)

def train_dataloader(self) -> DataLoader:
Expand Down
45 changes: 1 addition & 44 deletions pixi.lock
View file
Open in desktop

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

4 changes: 3 additions & 1 deletion pixi.toml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -36,6 +36,8 @@
# | H100 cuda13 | to test | to test | to test | to test |
# | B300 cuda12 | works | works | to test | to test |
# | B300 cuda13 | works | works | to test | to test |
# | GB10 cuda12 | n/a | n/a | works | works |
# | GB10 cuda13 | n/a | n/a | works | works |
# | CPU | works | n/a | to test | n/a |
# +---------------+------------+------------+------------+------------+
# * works = "works for me"
Expand Down Expand Up @@ -368,7 +370,7 @@ torch = { version = "*", index = "https://download.pytorch.org/whl/cu130" }
# 2.1) No aarch wheel for cuequivariance-ops-torch-cu13 <0.8
#

[feature.cuequivariance-cuda12.pypi-dependencies]
[feature.cuequivariance-cuda12.target.linux-64.pypi-dependencies]
cuequivariance = "<0.8"
cuequivariance-torch = "<0.8"
cuequivariance-ops-torch-cu12 = "<0.8"
Expand Down