10 Commits

Author	SHA1	Message	Date
Ross Wightman	162f49295e	Move naflex global pool into one fn that can be marked notrace	2025-05-24 14:06:12 -07:00
Ross Wightman	2ad75e8023	Fix issue w/ MAP attention mask and no patch_valid	2025-05-24 12:37:39 -07:00
Ross Wightman	d7d3538335	Add so400m model size for test, few tweaks.	2025-05-23 18:29:30 -07:00
Ross Wightman	ee27b73da4	Further pos embed tweaks, rejig model defs for testing	2025-04-28 09:15:11 -07:00
Ross Wightman	3dc90ed7a7	Add naflex loader support to validate.py, fix bug in naflex pos embed add, classic vit weight loading for naflex model	2025-04-25 16:00:54 -07:00
Ross Wightman	c527c37969	Optimizations for pos embed resize, merge different mask helper fns	2025-04-21 14:05:18 -07:00
Ross Wightman	ea728f67fa	Improve several typing issues for flex vit, can (almost) work with jit if we bash h,w key into an int or str	2025-04-14 11:01:56 -07:00
Ross Wightman	b4bb0f452a	Exclude embeds module and mask attn functions from tracing	2025-04-09 15:34:15 -07:00
Ross Wightman	13e0f3a4a3	Add loss scale arg, initial distributed loss scale. Maybe fix FX for the model.	2025-04-08 20:47:57 -07:00
Ross Wightman	0893f5d296	Initial NaFlex ViT model and training support	2025-04-07 21:27:10 -07:00