Early September Improvements + Fixes

timaeus-research / devinterp

Tools for studying developmental interpretability in neural networks.

77 stars 14 forks source link

Early September Improvements + Fixes #95

Closed wz-ml closed 2 months ago

wz-ml commented 2 months ago

Additions:

[x] Fixes issue #94.
[x] Multi-GPU support.
- [x] Multi-GPU support with one core per GPU.
- [x] Support with multiple cores per GPU
- [x] Verify sampling chains match
- [x] Multi-GPU test.
[x] Batch accumulation bugfix (default_nbeta)
- [x] Batch accumulation test
[x] Sampling optimizations
- [x] Automatic mixed precision
  - [x] AMP tests
[x] devinterp[tpu] torch_xla versioning dependency bugfix.

svwingerden commented 2 months ago

This looks great, thanks also for adding in some more comments / docstrings. Could you write a quick check that multi-GPU llc estimation works? doesn't need to converge to any known value, i'd just like to have some indication when someone inevitably breaks this in the future

wz-ml commented 2 months ago

Sure, I'll write a test!

wz-ml commented 2 months ago

Thanks for the feedback! Let me know if any more issues with this PR need to be resolved. Do our Github actions runner have more than 1 core?