Skip to content

working on oppy #1

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
afermg merged 1 commit into from
Jul 25, 2025
Merged

working on oppy #1

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
2 changes: 1 addition & 1 deletion dinov2/configs/train/vitl14_rxrx3_core.yaml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,7 @@ ibot:
separate_head: true
head_n_prototypes: 131072
train:
batch_size_per_gpu: 16
batch_size_per_gpu: 96
dataset_path: RXRX3_CORE
centering: sinkhorn_knopp
student:
Expand Down
16 changes: 13 additions & 3 deletions dinov2/train/ssl_meta_arch.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -345,12 +345,22 @@ def get_teacher_output():

return loss_dict

# def fsdp_synchronize_streams(self):
# if self.need_to_synchronize_fsdp_streams:
# torch.cuda.synchronize()
# self.student.dino_head._streams = (
# self.teacher.dino_head._streams
# ) = self.student.backbone._streams = self.teacher.backbone._streams
# self.need_to_synchronize_fsdp_streams = False

def fsdp_synchronize_streams(self):
if self.need_to_synchronize_fsdp_streams:
torch.cuda.synchronize()
self.student.dino_head._streams = (
self.teacher.dino_head._streams
) = self.student.backbone._streams = self.teacher.backbone._streams
for attr in {"_unshard_stream", "_post_backward_stream", "_pre_unshard_stream", "_all_reduce_stream", "_default_stream"}:
stream = getattr(self.teacher.backbone, attr)
setattr(self.student.dino_head, attr, stream)
setattr(self.teacher.dino_head, attr, stream)
setattr(self.student.backbone, attr, stream)
self.need_to_synchronize_fsdp_streams = False

def update_teacher(self, m):
Expand Down
8 changes: 4 additions & 4 deletions flake.lock
View file
Open in desktop

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

15 changes: 12 additions & 3 deletions flake.nix
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
{
inputs = {
nixpkgs.url = "github:NixOS/nixpkgs/nixos-unstable";
nixpkgs.url = "github:NixOS/nixpkgs/nixos-24.11";
nixpkgs_master.url = "github:NixOS/nixpkgs/master";
systems.url = "github:nix-systems/default";
flake-utils.url = "github:numtide/flake-utils";
Expand Down Expand Up @@ -42,7 +42,7 @@
devShells = {
default = let
python_with_pkgs = (
python3.withPackages (pp: [
python312.withPackages (pp: [
(inputs.nahual-flake.packages.${system}.nahual)
packages.dinov2
])
Expand All @@ -61,8 +61,17 @@
unset SOURCE_DATE_EPOCH
'';
shellHook = ''
export CUDA_PATH=${pkgs.cudaPackages.cudatoolkit}
export LD_LIBRARY_PATH=${pkgs.cudaPackages.cudatoolkit}/lib:${pkgs.cudaPackages.cudnn}/lib:$LD_LIBRARY_PATH
export NVCC_APPEND_FLAGS="-Xcompiler -fno-PIC"
export TORCH_CUDA_ARCH_LIST="6.0;6.1;7.0;7.5;8.0;8.6"
export CUDA_NVCC_FLAGS="-O2 -Xcompiler -fno-PIC"
runHook venvShellHook
export PYTHONPATH=${python_with_pkgs}/${python_with_pkgs.sitePackages}:$PYTHONPATH
# export PYTHONPATH=${python_with_pkgs}/${python_with_pkgs.sitePackages}:$PYTHONPATH
# Set PYTHONPATH to only include the Nix packages, excluding current directory
export PYTHONPATH=${python_with_pkgs}/${python_with_pkgs.sitePackages}
# Ensure current directory is not in Python path
export PYTHONDONTWRITEBYTECODE=1
'';
};
};
Expand Down
4 changes: 2 additions & 2 deletions nix/dinov2.nix
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -13,7 +13,7 @@
fvcore,
iopath,
xformers,
#submitit,
submitit,
# Extras
mmcv,
# Server,
Expand Down Expand Up @@ -44,7 +44,7 @@ buildPythonPackage {
fvcore
iopath
xformers
#submitit
submitit
# extras
# mmcv
# Server
Expand Down